protein bert uniref90 dataset about progen HOT 4 OPEN

import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.dataset as ds
import pandas as pd

from torch.utils.data import IterableDataset
from torch.utils.data import get_worker_info
from torch.multiprocessing import Queue

class IterableManualParquetDataset(IterableDataset):
    def __init__(self, path, process_func, batch_size=64):
        super().__init__()
        self.dataset = ds.dataset(path)
        self.batch_size = batch_size
        self.process_func = process_func

    def __iter__(self):
        worker_info = get_worker_info()

        # Only divide up batches when using multiple worker processes
        if worker_info != None:
            batches = list(self.dataset.to_batches(batch_size=self.batch_size))
            worker_load = len(batches) // worker_info.num_workers

            # If more workers than batches exist, some won't be used
            if worker_load == 0:
                if worker_info.id < len(batches): self.batches = [batches[worker_info.id]]
                else: return
            else:
                start = worker_load * worker_info.id
                end = min(start + worker_load, len(batches))
                self.batches = batches[start:end]
        else: self.batches = self.dataset.to_batches(batch_size=self.batch_size)

        # Process and yield each batch
        for batch in self.batches:
            batch = batch.to_pydict()
            batch.update(self.process_func(batch))

            yield batch
            
a = IterableManualParquetDataset("uniref90_parquet/uniref90_with_annotations", lambda x:x, batch_size=64)
u = next(iter(a))

import dask.dataframe as dd
from dask.distributed import Client
from pyspark.sql import SparkSession
from pyspark.sql.functions import col
if __name__ == '__main__':
    client = Client()
    def read_db_save_to_parquet():
            daskDF = dd.read_sql_table('protein_annotations', "sqlite:///uniref_proteins_and_annotations.db", index_col='index')
            daskDF.to_parquet('the-parquet', engine='pyarrow')

    def read_fasta_save_to_parquet():
        fasta_df = dd.read_csv("uniref90.fasta",lineterminator=">", sep="$", header=None)
        def process_fasta_row(row):
            lines = row[0].split("\n")
            id = lines[0].split(" ")[-1].split("=")[-1]
            seq = "".join(lines[1:])
            return (id, seq)
        new_df = fasta_df.apply(process_fasta_row, axis=1,result_type='expand', meta={0: str, 1: str})
        new_df.columns = ['uniprot_id', 'seq']
        new_df.to_parquet('fasta_parquet', engine='pyarrow')

    def spark_join():
        # executor and driver not sure if needed
        spark = SparkSession.builder \
            .config("spark.executor.cores", "16") \
            .config("spark.executor.memory", "16G") \
            .config("spark.driver.cores", "16") \
            .config("spark.driver.memory", "16G") \
            .master("local[16]").appName('spark-merge').getOrCreate() 
        a = spark.read.parquet("the-parquet")
        b = spark.read.parquet("fasta_parquet")
        c = a.join(b, col("uniprot_name") == col("uniprot_id"))
        c.write.mode("overwrite").parquet("uniref90_with_annotations")

    def merge_without_indexing():
        db_df = dd.read_parquet('the-parquet', engine='pyarrow') 
        fasta_index_df = dd.read_parquet('fasta_parquet', engine='pyarrow')
        merged = db_df.merge(fasta_index_df, left_on=("uniprot_name"), right_on=("uniprot_id"))
        merged.to_parquet('merged', engine='pyarrow')

import dask.dataframe as dd
from dask.distributed import Client
client = Client()
merged = dd.read_parquet('uniref90_with_annotations', engine='pyarrow') 
merged.head()