pysz 0.1.0__tar.gz

pysz-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Jinyang Zhang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pysz-0.1.0/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.1
+Name: pysz
+Version: 0.1.0
+Summary: Package for reading and writing svb-zstd compressed data
+Author: Jinyang Zhang
+Author-email: zhangjinyang@biols.ac.cn
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Requires-Dist: numpy (>=1.24.2,<2.0.0)
+Requires-Dist: pandas (>=1.5.3,<2.0.0)
+Requires-Dist: pystreamvbyte
+Requires-Dist: zstandard (==0.19.0)

pysz-0.1.0/pyproject.toml

@@ -0,0 +1,50 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = [
+    "setuptools",
+    "wheel",
+]
+
+[project]
+name = "pysz"
+version = "0.1.0"
+authors = [
+    {name = "Jinyang Zhang", email = "zhangjinyang@biols.ac.cn"},
+]
+description = "Package for reading and writing svb-zstd compressed data"
+readme = {file = "README.md", content-type = "text/markdown"}
+requires-python = "==3.10"
+keywords = ["svb", "zstd"]
+license = {file = "LICENSE"}
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3"
+]
+
+dependencies = [
+    "numpy==1.24.2",
+    "pandas==1.5.3",
+    "zstandard==0.19.0",
+    "pystreamvbyte @ git+https://github.com/Kevinzjy/pystreamvbyte@50eaf184b734fd7d64d18fd9b0847fb210e05b75#egg=pystreamvbyte",
+]
+
+[tool.poetry]
+name = "pysz"
+version = "0.1.0"
+description = "Package for reading and writing svb-zstd compressed data"
+authors = ["Jinyang Zhang <zhangjinyang@biols.ac.cn>"]
+
+
+[tool.poetry.dependencies]
+python = "^3.10"
+pandas = "^1.5.3"
+numpy = "^1.24.2"
+zstandard = "0.19.0"
+pystreamvbyte = "git+https://github.com/Kevinzjy/pystreamvbyte@50eaf184b734fd7d64d18fd9b0847fb210e05b75#egg=pystreamvbyte"
+
+[tool.poetry.dev-dependencies]
+pytest = "^7.2.2"
+
+#[tool.poetry.plugins."scripts"]
+#example_etl = "pysz.api:check_szmain"

pysz-0.1.0/setup.py

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+from setuptools import setup
+
+package_dir = \
+{'': 'src'}
+
+packages = \
+['pysz']
+
+package_data = \
+{'': ['*'], 'pysz': ['.pytest_cache/*', '.pytest_cache/v/cache/*']}
+
+install_requires = \
+['numpy>=1.24.2,<2.0.0',
+ 'pandas>=1.5.3,<2.0.0',
+ 'pystreamvbyte',
+ 'zstandard==0.19.0']
+
+setup_kwargs = {
+    'name': 'pysz',
+    'version': '0.1.0',
+    'description': 'Package for reading and writing svb-zstd compressed data',
+    'long_description': None,
+    'author': 'Jinyang Zhang',
+    'author_email': 'zhangjinyang@biols.ac.cn',
+    'maintainer': None,
+    'maintainer_email': None,
+    'url': None,
+    'package_dir': package_dir,
+    'packages': packages,
+    'package_data': package_data,
+    'install_requires': install_requires,
+    'python_requires': '>=3.10,<4.0',
+}
+
+
+setup(**setup_kwargs)

pysz-0.1.0/src/pysz/.pytest_cache/.gitignore

@@ -0,0 +1,2 @@
+# Created by pytest automatically.
+*

pysz-0.1.0/src/pysz/.pytest_cache/CACHEDIR.TAG

@@ -0,0 +1,4 @@
+Signature: 8a477f597d28d172789f06886806bc55
+# This file is a cache directory tag created by pytest.
+# For information about cache directory tags, see:
+#	https://bford.info/cachedir/spec.html

pysz-0.1.0/src/pysz/.pytest_cache/README.md

@@ -0,0 +1,8 @@
+# pytest cache directory #
+
+This directory contains data from the pytest's cache plugin,
+which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
+
+**Do not** commit this to version control.
+
+See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

pysz-0.1.0/src/pysz/.pytest_cache/v/cache/lastfailed

@@ -0,0 +1,3 @@
+{
+  "api.py::test_read": true
+}

pysz-0.1.0/src/pysz/.pytest_cache/v/cache/nodeids

@@ -0,0 +1,3 @@
+[
+  "api.py::test_read"
+]

pysz-0.1.0/src/pysz/.pytest_cache/v/cache/stepwise

@@ -0,0 +1 @@
+[]

pysz-0.1.0/src/pysz/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

pysz-0.1.0/src/pysz/api.py

@@ -0,0 +1,325 @@
+import json
+import numpy as np
+import pandas as pd
+from pathlib import Path
+from collections import OrderedDict, namedtuple
+from zstandard import ZstdDecompressor
+
+from pysz import __version__
+from pysz.compression import svb_decode, svb_encode, zstd_decode, zstd_encode, str_decode, str_encode
+from pysz.utils import mkdir, assert_file_exists, assert_dir_exists
+
+from multiprocessing import Manager, Process
+
+# Default attributes & datasets for chunk data
+Attributes = [
+    ('ID', str),
+    ('Offset', np.int32),
+    ('Raw_unit', np.uint32),
+]
+Datasets = [
+    ('Raw', np.uint32),
+    ('Fastq', str),
+    ('Move', np.uint16),
+]
+
+# For dtype conversion
+Dtypes_names = {
+    str: "S",
+    np.int16: "I16", np.int32: "I32",
+    np.uint16: "U16", np.uint32: "U32",
+    np.float16: "F16", np.float32: "F32",
+}
+Dtypes = {j: i for i, j in Dtypes_names.items()}
+
+# Dtypes that can be svb compressed
+Svb_dtypes = [np.int16, np.int32, np.uint16, np.uint32]
+
+
+class CompressedFile(object):
+    """
+    Class for parsing SVB-ZSTD compressed data
+    """
+
+    def __init__(self, data_dir, mode:str="r", header:list=None, attributes:list=None, datasets:list=None, overwrite:bool=False, n_threads=8, allow_multiprocessing=True):
+        """
+        Init function of CompressedFile Class
+
+        Args:
+            data_dir: str / Path,
+                path to SZ directory, required
+            mode: str,
+                "r" for read and "w" for writing sz data
+            header: [(key, value)]:
+                list of of keys and values in header, defaults to [(version, '0.0.1')]
+            attributes: [(attr_id, attr_dtype)],
+                list of ID and dtype for attributes, defaults to [('ID', str), ('Offset', np.int32), ('Raw_unit', np.uint32)]
+            datasets: [(dataset_id, dataset_dtype)],
+                list of ID and dtype for datasets, defaults to [('Raw', np.uint32), ('Fastq', np.uint32), ('Move', np.uint32)]
+            overwrite: bool,
+                whether overwrite existed sz data, defaults to False
+            n_threads: int,
+                number of threads for parallelized data compression, defaults to 8
+            allow_multiprocessing: bool,
+                set to True if want to access SZ data using multiprocessing, defaults to False
+        """
+        self.dir = data_dir if isinstance(data_dir, Path) else Path(data_dir)
+        self.idx_path = self.dir / "index"
+        self.dat_path = self.dir / "dat"
+        self.mode = mode
+
+        # For Record class
+        self.idx, self.header, self.attr, self.datasets = None, None, None, None
+
+        # For multiprocessing
+        self.threads = n_threads
+        self.q_in = Manager().Queue()
+        self.q_out = Manager().Queue()
+        self.pool = []
+        self.worker, self.decompressor, self.fh = None, None, None
+
+        # Get attributes and datasets
+        if self.mode == "r":
+            assert_dir_exists(data_dir)
+            assert_file_exists(self.idx_path)
+            assert_file_exists(self.dat_path)
+            self.load_header()
+        elif self.mode == 'w':
+            mkdir(data_dir, overwrite=overwrite)
+            self.save_header(header, attributes, datasets)
+        else:
+            raise KeyError("Mode should be either 'w' or 'r'.")
+
+        # Init Record class
+        keys = []
+        if self.attr is not None:
+            keys += list(self.attr)
+        if self.datasets is not None:
+            keys += list(self.datasets)
+        self.record = namedtuple("Record", keys)
+
+        # Init multiprocessing related variables
+        if self.mode == 'w':
+            self.worker = Process(target=self.writer)
+            self.worker.start()
+            for _ in np.arange(self.threads):
+                p = Process(target=self.encoder)
+                p.start()
+                self.pool.append(p)
+        else:
+            if allow_multiprocessing:
+                self.decompressor = ZstdDecompressor()
+                self.fh = open(self.dat_path, 'rb')
+
+    def save_header(self, header, attributes, datasets):
+        """
+        Save header, attributes and dataset information to the index file
+
+        Args:
+            header: [(key, value)]:
+                list of of keys and values in header, defaults to [(version, '0.0.1')]
+            attributes: [(attr_id, attr_dtype)],
+                list of ID and dtype for attributes, defaults to [('ID', str), ('Offset', np.int32), ('Raw_unit', np.uint32)]
+            datasets: [(dataset_id, dataset_dtype)],
+                list of ID and dtype for datasets, defaults to [('Raw', np.uint32), ('Fastq', np.uint32), ('Move', np.uint32)]
+        """
+        self.header =  {"version": __version__} if header is None else dict(header)
+        self.attr =  OrderedDict(Attributes) if attributes is None else OrderedDict(attributes)
+        self.datasets = OrderedDict(Datasets) if datasets is None else OrderedDict(datasets)
+
+        with open(self.idx_path, 'w') as out:
+            header_str = json.dumps(self.header)
+            out.write("#" + header_str + '\n')
+
+            cols = ["LENGTH", "OFFSET",] + \
+                [f"{i}:A:{Dtypes_names[j]}" for i,j in self.attr.items()] + \
+                [f"{i}:D:{Dtypes_names[j]}" for i,j in self.datasets.items()]
+            out.write('\t'.join(cols) + '\n')
+
+    def load_header(self):
+        """
+        Load header, attributes and datasets information for SZ file
+        """
+        attributes = []
+        datasets = []
+        index = []
+        with open(self.idx_path, 'r') as f:
+            # Get first row of header
+            header = json.loads(f.readline().rstrip().lstrip('#'))
+
+            # Get attributes and datasets
+            cols = f.readline().rstrip().split('\t')
+            for col in cols[2:]:
+                key_id, key_group, key_dtype = col.split(':')
+                if key_group == 'A':
+                    attributes.append((key_id, Dtypes[key_dtype]))
+                elif key_group == "D":
+                    datasets.append((key_id, Dtypes[key_dtype]))
+                else:
+                    raise KeyError(f"Unsupported data type: {key_group}")
+
+            # Get index dataframe
+            for line in f:
+                index.append(line.rstrip().split('\t'))
+
+        self.header = header
+        self.attr = OrderedDict(attributes)
+        self.datasets = OrderedDict(datasets)
+
+        # Convert into pandas dataframe and parser dtype
+        self.idx = pd.DataFrame(index)
+        self.idx.columns = ['LENGTH', 'OFFSET'] + list(self.attr) + list(self.datasets)
+        self.idx['LENGTH'] = self.idx['LENGTH'].astype(int)
+        self.idx['OFFSET'] = self.idx['OFFSET'].astype(int)
+
+    def put(self, *data):
+        """
+        Put new data into queue for compressing and saving
+
+        Args:
+            *data: data must be in specific order consistent to the attributes and datasets.
+                The default SZ contains the following attributes:
+
+                          ID: read ID, str
+                      Offset: Offset for raw current data, np.int32
+                    Raw_unit: Raw unit for raw current data, np.float32
+
+                And the following datasets:
+
+                        Raw: Raw current data points, np.uint32
+                      Fastq: Basecalled fastq, str
+                       Move: move tables, np.uint16
+        """
+        self.q_in.put(data)
+
+    def encode(self, data):
+        """
+        Encode input data into SVB-ZSTD compressed binary format
+        Args:
+            data: data acquired from queue
+
+        Returns:
+            encoded, byte str,
+                byte str of compressed datasets
+            idx, list,
+                list of items to be stored in index
+        """
+        record = self.record(*data)
+        dat_encoded = []
+        idx_encoded = []
+
+        # Store indices
+        for attr_id, attr_dtype in self.attr.items():
+            idx_encoded.append(str(attr_dtype(getattr(record, attr_id))))
+
+        # Compress datasets according to the dtypes
+        for dataset_id, dataset_dtype in self.datasets.items():
+            if dataset_dtype in Svb_dtypes:
+                d_bstr, d_size, _ = svb_encode(np.array(getattr(record, dataset_id)).astype(dataset_dtype))
+            elif dataset_dtype == str:
+                d_bstr = str_encode(dataset_dtype(getattr(record, dataset_id)))
+                d_size = ''
+            else:
+                d_bstr = np.array(getattr(record, dataset_id)).astype(dataset_dtype).tobytes()
+                d_size = ''
+
+            d_encoded = zstd_encode(d_bstr)
+            dat_encoded.append(d_encoded)
+            idx_encoded.append(f"{len(d_encoded)}:{d_size}")
+
+        # Convert encoded into simple binary format
+        encoded = b''.join(dat_encoded)
+
+        return encoded, idx_encoded
+
+    def encoder(self):
+        """
+        Process for compressing input data and pass to SZ writer
+        """
+        while True:
+            raw = self.q_in.get()
+            if raw is None:
+                break
+            dat_line, idx_line = self.encode(raw)
+            self.q_out.put((dat_line, idx_line))
+        self.q_out.put(None)
+
+    def writer(self):
+        """
+        Process for handling compressed data and save it to disk
+        """
+        n_alive = self.threads
+        cursor = 0
+        with open(self.idx_path, 'a') as idx_out, open(self.dat_path, 'wb') as dat_out:
+            while True:
+                # Wait for all encode workers to stop
+                if n_alive == 0:
+                    break
+                res = self.q_out.get()
+                if res is None:
+                    n_alive -= 1
+                    continue
+
+                # Record cursor position everytime
+                dat_line, idx_line = res
+                dat_out.write(dat_line)
+                shift, length = cursor, len(dat_line)
+                cursor += length
+                idx_line = [str(length), str(shift), ] + idx_line
+                idx_out.write('\t'.join(idx_line) + '\n')
+
+    def get(self, idx):
+        """
+        Get specified index from compressed SZ data
+
+        Args:
+            idx: int / list
+                numeric int or list of integer index, items will be acquired using 'df.iloc';
+
+        Returns:
+            list of Record (namedTuple), containing attributes and datasets information specified in the SZ header
+
+        """
+        reads = []
+        rows = self.idx.iloc[[idx]] if isinstance(idx, int) else self.idx.iloc[idx]
+        rows = rows.sort_values(by='OFFSET')
+
+        zdc = ZstdDecompressor() if self.decompressor is None else self.decompressor
+        f = open(self.dat_path, 'rb') if self.fh is None else self.fh
+
+        for _, row in rows.iterrows():
+            offset, length = row['OFFSET'], row['LENGTH']
+            attr = row[2:2+len(self.attr)].tolist()
+
+            f.seek(offset)
+            encoded = f.read(length)
+
+            datasets = []
+            p = 0
+            for d_name, d_col in zip(self.datasets, row[2+len(self.attr):]):
+                d_size, d_shape = d_col.split(':')
+                bstr = zstd_decode(encoded[p:p+int(d_size)], zdc)
+                p += int(d_size)
+
+                d_dtype = self.datasets[d_name]
+                if d_dtype in Svb_dtypes:
+                    d_decoded = svb_decode(bstr, int(d_shape), dtype=d_dtype)
+                elif d_dtype == str:
+                    d_decoded = str_decode(bstr)
+                else:
+                    d_decoded = np.frombuffer(bstr, dtype=d_dtype)
+                datasets.append(d_decoded)
+            record = attr + datasets
+            reads.append(self.record(*record))
+        return reads
+
+    def close(self):
+        """
+        Functions for closing and wanting for process to end.
+        """
+        if self.mode == 'w':
+            for _ in np.arange(self.threads):
+                self.q_in.put(None)
+            for p in self.pool:
+                p.join()
+            self.worker.join()

pysz-0.1.0/src/pysz/compression.py

@@ -0,0 +1,108 @@
+import numpy as np
+import streamvbyte
+from zstandard import ZstdDecompressor, ZstdCompressor
+
+
+def str_to_bytes(bytes_or_str):
+    """
+    Convert str to byte string
+    :param bytes_or_str:
+    :return:
+    """
+    if isinstance(bytes_or_str, str):
+        value = bytes_or_str.encode('utf-8')
+    else:
+        value = bytes_or_str
+    return value
+
+
+def bytes_to_str(bytes_or_str):
+    """
+    Convert byte string to string
+    :param bytes_or_str:
+    :return:
+    """
+    if isinstance(bytes_or_str, bytes):
+        value = bytes_or_str.decode('utf-8')
+    else:
+        value = bytes_or_str
+    return value
+
+
+def str_encode(data):
+    """
+    Convert str to byte
+    :param data:
+    :return:
+    """
+    return str_to_bytes(data)
+
+
+def str_decode(data):
+    """
+    Convert byte to str
+    :param data:
+    :return:
+    """
+    return bytes_to_str(data)
+
+
+def is_gzipped(file_name):
+    """
+    Examine whether a file is gzipped.
+    Reference: https://stackoverflow.com/questions/3703276/how-to-tell-if-a-file-is-gzip-compressed
+    """
+    with open(file_name, 'rb') as f:
+        return f.read(2) == b'\x1f\x8b'
+
+
+def zstd_encode(bstr, cctx=None):
+    """
+    Compress byte string using Zstd library
+    :param data: byte str
+        data for compression
+    :param cctx: zstandard.ZstdCompressor
+        instance of zstd compressor
+    :return: byte str
+        compressed data
+    """
+    _cctx = ZstdCompressor() if cctx is None else cctx
+    compressed = _cctx.compress(bstr)
+    return compressed
+
+
+def zstd_decode(compressed, zdc=None):
+    """
+    Decompress zstd compressed data
+    :param compressed:
+    :param zdc:
+    :return:
+    """
+    _zdc = ZstdDecompressor() if zdc is None else zdc
+    decompressed = _zdc.decompress(compressed)
+    return decompressed
+
+
+def svb_encode(data):
+    """
+    Compress data using streamvbyte algorithm, int16, uint16, int32, uint32 is supported
+    :param data:
+    :return:
+    """
+    if data.dtype not in [np.uint16, np.int16, np.int32, np.uint32]:
+        raise TypeError("Streamvbyte only support int16/uint16/int32/uint32 formats!")
+    compressed = streamvbyte.encode(data).tobytes()
+    return compressed, data.shape[0], data.dtype
+
+
+def svb_decode(data, size, dtype=np.int32):
+    """
+    Decompress streamvbyte encoded data
+    :param data:
+    :param size:
+    :param dtype:
+    :return:
+    """
+    decoded = np.frombuffer(data, dtype=np.uint8)
+    decompressed = streamvbyte.decode(decoded, size, dtype=dtype)
+    return decompressed

pysz-0.1.0/src/pysz/utils.py

@@ -0,0 +1,35 @@
+import sys
+import warnings
+from pathlib import Path
+
+def mkdir(dir_name, overwrite=False):
+    dir_path = dir_name if isinstance(dir_name, Path) else Path(dir_name)
+    if dir_path.exists():
+        if dir_path.is_dir():
+            if overwrite is True:
+                sys.stderr.write(f'Overwriting existed {dir_name}\n')
+            else:
+                raise OSError(f'{dir_name} already exists, please use --overwrite to force overwrite')
+        else:
+            raise OSError(f'{dir_name} conflict with existing files!')
+    else:
+        dir_path.mkdir()
+    return dir_path
+
+
+def assert_dir_exists(dir_name):
+    dir_path = dir_name if isinstance(dir_name, Path) else Path(dir_name)
+    if dir_path.exists() and dir_path.is_dir():
+        return dir_path
+    raise OSError(f"Directory not found: {dir_path}.")
+
+
+def assert_file_exists(file_name):
+    file_path = file_name if isinstance(file_name, Path) else Path(file_name)
+    if file_path.exists() and file_path.is_file():
+        return file_path
+    raise OSError(f"File not found: {file_path}.")
+
+
+def error_callback(e):
+    sys.stderr.write(f"Error in worker process: {e.__cause__}: {e}")