Package: pysz
pyszPython interface for writing and reading svb-zstd compressed dataInstallationDependenciesPython >= 3.8Numpy >= 1.24.2Pandas >= 2.0.0Zstandard >= 0.20.0Install using pip# First, install a customized version of pystreamvbytepipinstallgit+https://github.com/kevinzjy/pystreamvbyte.git# Install pysz nowpipinstallpyszInstall from source# First, install a customized version of pystreamvbytepipinstallgit+https://github.com/kevinzjy/pystreamvbyte.git# Install pysz from sourcegitclone--recursivehttps://github.com/kevinzjy/pyszcdpysz
pipinstall.UsageCreate & Write SZ fileSupported data types：strnp.int16 / np.int32np.uint16 / np.uint32np.float16 / np.float32 / np.float64 / np.float128Example for creating SZ fileimportnumpyasnpfrompysz.apiimportCompressedFileheader=[('version','0.0.1'),('date','2023-04-03')]attr=[('ID',str),('Offset',np.int32),('Raw_unit',np.float32)]datasets=[('Raw',np.uint32),('Fastq',str),('Move',np.uint16)]# Create new SZ filesz=CompressedFile("/tmp/test_sz",mode="w",header=header,attributes=attr,datasets=datasets,overwrite=True,n_threads=8)# Save data in single-read modeforiinrange(10000):sz.put(f"read_{i}",# ID0,# Offsetnp.random.rand(),# Raw_unitnp.random.randint(70,150,4000),# Raw''.join(np.random.choice(['A','T','C','G'],450)),# Fastqnp.random.randint(0,1,4000),# Move)# Save data in chunk modeforiinrange(100):chunk=[]forjinrange(100):chunk.append((f"read_{i}_{j}",0,np.random.rand(),np.random.randint(70,150,4000),''.join(np.random.choice(['A','T','C','G'],450)),np.random.randint(0,1,4000),))sz.put_chunk(chunk)# Remember to call it to ensure everything finished successfullysz.close()Example for creating SZ file with multiprocessingNote: for creating SZ file with multiprocessing, pass theCompressedFile.q_in(a multiprocessing.Manager().Queue() instance)
as parameter to parallelized function. Then usequeue.put((True, chunk))andqueue.put((True, read))for saving reads
or chunks directly.importnumpyasnpfrompysz.apiimportCompressedFilefrommultiprocessingimportPooldefwriter(queue,chunk_id):# Write in single-read modeforiinrange(10):read=(f"read_{chunk_id}_{i}",0,np.random.rand(),np.random.randint(70,150,4000),''.join(np.random.choice(['A','T','C','G'],450)),np.random.randint(0,1,4000),np.random.random(4000),)queue.put((False,read))# Write in chunk modechunk=[]foriinrange(10):chunk.append((f"chunk_{chunk_id}_{i}",0,np.random.rand(),np.random.randint(70,150,4000),''.join(np.random.choice(['A','T','C','G'],450)),np.random.randint(0,1,4000),np.random.random(4000),))queue.put((True,chunk))header=[('version','0.0.1'),('date','2023-04-03')]attr=[('ID',str),('Offset',np.int32),('Raw_unit',np.float32)]datasets=[('Raw',np.uint32),('Fastq',str),('Move',np.uint16)]# Create new SZ filesz=CompressedFile("/tmp/test_sz",mode="w",header=header,attributes=attr,datasets=datasets,overwrite=True,n_threads=8)# Use pool for multiprocessingpool=Pool(8)jobs=[]foriinrange(10):jobs.append(pool.apply_async(writer,(sz.q_in,i,)))pool.close()pool.join()# Remember to call it to ensure everything finished successfullysz.close()Examples for reading SZ fileimportnumpyasnpfrompysz.apiimportCompressedFilesz=CompressedFile("/tmp/test_sz",mode="r",allow_multiprocessing=True,)# Get index informationprint(sz.idx.head())# Get total read numberprint(f"Loaded{sz.idx.shape[0]}reads")# Get the first readread=sz.get(0)print(read.ID,read.Offset,read.Raw_unit,read.Raw,read.Fastq,read.Move)# Get first 10 readsreads=sz.get(np.arange(10))# Filter some readsidx=sz.idx.index[sz.idx['ID'].isin(['read_0','read_1'])]reads=sz.get(idx)Example for reading SZ file chunk by chunk using multiprocessingNote: for reading SZ file with multiprocessing, pass the chunked index as parameter,
and init CompressedFile instance withallow_multiprocessing=Trueseparately in each process.importnumpyasnpfrompysz.apiimportCompressedFilefrompysz.utilsimporterror_callbackfrommultiprocessingimportPooldefgrouper(iterable,n,fillvalue=None):"""Collect data info fixed-length chunks or blocksgrouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"""fromitertoolsimportzip_longestargs=[iter(iterable)]*nreturnzip_longest(*args,fillvalue=fillvalue)defread_from_sz(file_name,idx):"""Read from SZ inside each process"""sz=CompressedFile(file_name,mode='r',allow_multiprocessing=True)reads=sz.get(idx)sz.close()returnreads# Init pooln_threads=16chunk_size=500# Get all reads stored in SZ filefile_name='/tmp/test_sz'sz=CompressedFile(file_name,mode='r')index=sz.idxsz.close()# Add process to Poolpool=Pool(n_threads)jobs=[]forxingrouper(index.index,chunk_size):chunk=[iforiinxifiisnotNone]jobs.append(pool.apply_async(read_from_sz,(file_name,chunk,),error_callback=error_callback))pool.close()# Get outputforjobinjobs:ret=job.get()# Process your data herepool.join()Please refer totests/test_api.pyfor detailed usage.TODOAdd support for zfp for lossy float compressionRefactor the naive solution for using multiprocessing for writing SZ.