zip_csv_data_source.py

from pathlib import Path

from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
from pyspark.sql.types import StructType


class RangePartition(InputPartition):
 def __init__(self, start, end):
 self.start = start
 self.end = end


class ZipDataSourceReader(DataSourceReader):

 def __init__(self, schema, options):
 self.schema: StructType = schema
 self.options = options
 self.path = self.options.get("path", None)
 self.numPartitions = int(self.options.get("numPartitions", 2))
 print(options)

 def partitions(self):
 return [RangePartition(0, 1000) for i in range(self.numPartitions)]

 def read(self, partition):
 # Library imports must be within the method.
 from zipfile import ZipFile

 print(partition)
 try:
 p = Path(self.path)
 if not p.exists():
 print(f"not exists {p}")
 return
 if p.is_dir():
 # a folder full of zips
 # beaware of .glob() at extreme
 for file in Path(self.path).glob("**/*.zip"):
 print(f"{file}")
 with ZipFile(file, "r") as zipFile:
 for name in zipFile.namelist():
 with zipFile.open(name, "r") as zipfile:
 for line in zipfile:
 yield [f"{file}, {name} {line.decode('utf-8')}"]
 else:
 # single zip file
 with ZipFile(file, "r") as zipFile:
 for name in zipFile.namelist():
 with zipFile.open(name, "r") as zipfile:
 for line in zipfile:
 print(type(line))
 yield [f"{file}, {name} {line.decode('utf-8')}"]
 except Exception as e:
 print(e)


from pyspark.sql.datasource import DataSource, DataSourceReader
from pyspark.sql.types import StructType


class ZipDataSource(DataSource):
 """
 An example data source for batch query using the `Zipr` library.
 """

 @classmethod
 def name(cls):
 return "Zip"

 def schema(self):
 return "line string"

 def reader(self, schema: StructType):
 return ZipDataSourceReader(schema, self.options)


from pyspark.sql import SparkSession
from pyspark.sql.functions import col


def create_spark_session(app_name="PySpark CSV Zips Datasource Tester"):
 """
 Creates and returns a Spark Session
 """
 return (
 SparkSession.builder.appName(app_name)
 .master("local[*]")
 .config("spark.memory.offHeap.enabled", "true")
 .config("spark.memory.offHeap.size", "1g")
 .getOrCreate()
 )


if __name__ == "__main__":
 print("Hello")

 spark = create_spark_session()
 spark.dataSource.register(ZipDataSource)
 # TODO: Fix re-create CSV data
 # df = spark.read.format("Zip").load("./zips/x.zip").limit(3)
 # print(df.collect())

 df = spark.read.format("Zip").load("./zips")
 df.write.format("text").mode("overwrite").save("./saves")