<commit_before>from enum import Enum
import json
import os
from shutil import copyfile
import tempfile
import zlib

import pandas as pd
import requests
from six import iteritems

try:
    import fastparquet
except ImportError:
    fastparquet = None

try:
    import pyarrow as pa
    from pyarrow import parquet
except ImportError:
    pa = None

try:
    from pyspark.sql import SparkSession
except ImportError:
    SparkSession = None

from .const import TargetType
from .core import (decode_node, encode_node, hash_contents,
                   FileNode, RootNode, GroupNode, TableNode,
                   PackageFormat)
from .hashing import digest_file

ZLIB_LEVEL = 2
ZLIB_METHOD = zlib.DEFLATED  # The only supported one.
ZLIB_WBITS = zlib.MAX_WBITS | 16  # Add a gzip header and checksum.
CHUNK_SIZE = 4096

class ParquetLib(Enum):
    ARROW = 'pyarrow'
    FASTPARQUET = 'fastparquet'
    SPARK = 'pyspark'


class PackageException(Exception):
    """
    Exception class for Package handling
    """
    pass


class Package(object):
    BUILD_DIR = 'build'
    OBJ_DIR = 'objs'
    TMP_OBJ_DIR = 'objs/tmp'
    DF_NAME = 'df'

    __parquet_lib = None

    @classmethod
    def get_parquet_lib(cls):
        if not cls.__parquet_lib:
            parq_env = os.environ.get('QUILT_PARQUET_LIBRARY')
            if parq_env:
                cls.__parquet_lib = ParquetLib(parq_env)
            else:
                if SparkSession is not None:
                    cls.__parquet_lib = ParquetLib.SPARK
                elif pa is not None:
                    cls.__parquet_lib = ParquetLib.ARROW
                elif fastparquet is not None:
                    cls.__parquet_lib = ParquetLib.FASTPARQUET
                else:
                    msg = "One of the following libraries is requried to read"
                    msg += " Parquet packages: %s" % [l.value for l in ParquetLib]
                    raise PackageException(msg)
        return cls.__parquet_lib

    @classmethod
    def reset_parquet_lib(cls):
        cls.__parquet_lib = None

    def __init__(self, user, package, path, pkg_dir):
        self._user = user
        self._package = package
        self._pkg_dir = pkg_dir
        self._path = path

    def file(self, hash_list):
        """
        Returns the path to an object file that matches the given hash.
        """
        assert isinstance(hash_list, list)
        assert len(hash_list) == 1, "File objects must be contained in one file."
        filehash = hash_list[0]
        return self._object_path(filehash)

    def _read_hdf5(self, hash_list):
        assert len(hash_list) == 1, "Multi-file DFs not supported in HDF5."
        filehash = hash_list[0]
        with pd.HDFStore(self._object_path(filehash), 'r') as store:
            return store.get(self.DF_NAME)

    def _read_parquet_arrow(self, hash_list):
        if pa is None:
            raise PackageException("Module pyarrow is required for ArrowPackage.")

        assert len(hash_list) == 1, "Multi-file DFs not supported for Arrow Packages (yet)."
        filehash = hash_list[0]

        nt = 8
        fpath = self._object_path(filehash)
        table = parquet.read_table(fpath, nthreads=nt)
        df = table.to_pandas()
        return df

    def _read_parquet_fastparquet(self, hash_list):
        assert len(hash_list) == 1, "Multi-file DFs not supported yet."
        filehash = hash_list[0]
        pfile = fastparquet.ParquetFile(self._object_path(filehash))
        return pfile.to_pandas()

    def _read_parquet_spark(self, hash_list):
        if SparkSession is None:
            raise PackageException("Module SparkSession from pyspark.sql is required for " +
                                   "SparkPackage.")

        spark = SparkSession.builder.getOrCreate()
        assert len(hash_list) == 1, "Multi-file DFs not supported yet."
        filehash = hash_list[0]
        df = spark.read.parquet(self._object_path(filehash))
        return df

    def _dataframe(self, hash_list, pkgformat):
        """
        Creates a DataFrame from a set of objects (identified by hashes).
        """
        enumformat = PackageFormat(pkgformat)
        if enumformat is PackageFormat.HDF5:
            return self._read_hdf5(hash_list)
        elif enumformat is PackageFormat.PARQUET:
            parqlib = self.get_parquet_lib()
            if parqlib is ParquetLib.SPARK:
                return self._read_parquet_spark(hash_list)
            elif parqlib is ParquetLib.ARROW:
                return self._read_parquet_arrow(hash_list)
            elif parqlib is ParquetLib.FASTPARQUET:
                return self._read_parquet_fastparquet(hash_list)
            else:
                assert False, "Unimplemented Parquet Library %s" % parqlib
        else:
            assert False, "Unimplemented package format: %s" % enumformat

    def save_df(self, df, name, path, ext, target):
        """
        Save a DataFrame to the store.
        """
        enumformat = PackageFormat(self.get_contents().format)
        buildfile = name.lstrip('/').replace('/', '.')
        storepath = self._temporary_object_path(buildfile)

        # Serialize DataFrame to chosen format
        if enumformat is PackageFormat.HDF5:
            with pd.HDFStore(storepath, mode='w') as store:
                store[self.DF_NAME] = df<commit_msg>Monkey-patch pandas to force the use of old pickle format (#61)<commit_after>from enum import Enum
import json
import os
from shutil import copyfile
import tempfile
import zlib

import pandas as pd
import requests
from six import iteritems

try:
    import fastparquet
except ImportError:
    fastparquet = None

try:
    import pyarrow as pa
    from pyarrow import parquet
except ImportError:
    pa = None

try:
    from pyspark.sql import SparkSession
except ImportError:
    SparkSession = None

from .const import TargetType
from .core import (decode_node, encode_node, hash_contents,
                   FileNode, RootNode, GroupNode, TableNode,
                   PackageFormat)
from .hashing import digest_file

ZLIB_LEVEL = 2
ZLIB_METHOD = zlib.DEFLATED  # The only supported one.
ZLIB_WBITS = zlib.MAX_WBITS | 16  # Add a gzip header and checksum.
CHUNK_SIZE = 4096

class ParquetLib(Enum):
    ARROW = 'pyarrow'
    FASTPARQUET = 'fastparquet'
    SPARK = 'pyspark'


class PackageException(Exception):
    """
    Exception class for Package handling
    """
    pass


class Package(object):
    BUILD_DIR = 'build'
    OBJ_DIR = 'objs'
    TMP_OBJ_DIR = 'objs/tmp'
    DF_NAME = 'df'

    __parquet_lib = None

    @classmethod
    def get_parquet_lib(cls):
        if not cls.__parquet_lib:
            parq_env = os.environ.get('QUILT_PARQUET_LIBRARY')
            if parq_env:
                cls.__parquet_lib = ParquetLib(parq_env)
            else:
                if SparkSession is not None:
                    cls.__parquet_lib = ParquetLib.SPARK
                elif pa is not None:
                    cls.__parquet_lib = ParquetLib.ARROW
                elif fastparquet is not None:
                    cls.__parquet_lib = ParquetLib.FASTPARQUET
                else:
                    msg = "One of the following libraries is requried to read"
                    msg += " Parquet packages: %s" % [l.value for l in ParquetLib]
                    raise PackageException(msg)
        return cls.__parquet_lib

    @classmethod
    def reset_parquet_lib(cls):
        cls.__parquet_lib = None

    def __init__(self, user, package, path, pkg_dir):
        self._user = user
        self._package = package
        self._pkg_dir = pkg_dir
        self._path = path

    def file(self, hash_list):
        """
        Returns the path to an object file that matches the given hash.
        """
        assert isinstance(hash_list, list)
        assert len(hash_list) == 1, "File objects must be contained in one file."
        filehash = hash_list[0]
        return self._object_path(filehash)

    def _read_hdf5(self, hash_list):
        assert len(hash_list) == 1, "Multi-file DFs not supported in HDF5."
        filehash = hash_list[0]
        with pd.HDFStore(self._object_path(filehash), 'r') as store:
            return store.get(self.DF_NAME)

    def _read_parquet_arrow(self, hash_list):
        if pa is None:
            raise PackageException("Module pyarrow is required for ArrowPackage.")

        assert len(hash_list) == 1, "Multi-file DFs not supported for Arrow Packages (yet)."
        filehash = hash_list[0]

        nt = 8
        fpath = self._object_path(filehash)
        table = parquet.read_table(fpath, nthreads=nt)
        df = table.to_pandas()
        return df

    def _read_parquet_fastparquet(self, hash_list):
        assert len(hash_list) == 1, "Multi-file DFs not supported yet."
        filehash = hash_list[0]
        pfile = fastparquet.ParquetFile(self._object_path(filehash))
        return pfile.to_pandas()

    def _read_parquet_spark(self, hash_list):
        if SparkSession is None:
            raise PackageException("Module SparkSession from pyspark.sql is required for " +
                                   "SparkPackage.")

        spark = SparkSession.builder.getOrCreate()
        assert len(hash_list) == 1, "Multi-file DFs not supported yet."
        filehash = hash_list[0]
        df = spark.read.parquet(self._object_path(filehash))
        return df

    def _dataframe(self, hash_list, pkgformat):
        """
        Creates a DataFrame from a set of objects (identified by hashes).
        """
        enumformat = PackageFormat(pkgformat)
        if enumformat is PackageFormat.HDF5:
            return self._read_hdf5(hash_list)
        elif enumformat is PackageFormat.PARQUET:
            parqlib = self.get_parquet_lib()
            if parqlib is ParquetLib.SPARK:
                return self._read_parquet_spark(hash_list)
            elif parqlib is ParquetLib.ARROW:
                return self._read_parquet_arrow(hash_list)
            elif parqlib is ParquetLib.FASTPARQUET:
                return self._read_parquet_fastparquet(hash_list)
            else:
                assert False, "Unimplemented Parquet Library %s" % parqlib
        else:
            assert False, "Unimplemented package format: %s" % enumformat

    def save_df(self, df, name, path, ext, target):
        """
        Save a DataFrame to the store.
        """
        enumformat = PackageFormat(self.get_contents().format)
        buildfile = name.lstrip('/').replace('/', '.')
        storepath = self._temporary_object_path(buildfile)

        # Serialize DataFrame to chosen format
        if enumformat is PackageFormat.HDF5:
            # HACK: Force the use of old pickle to ensure Python 2/3 compatibility.
            from pandas.compat import cPickle
            old_protocol = cPickle.HIGHEST_PROTOCOL
            try:
                cPickle.HIGHEST_PROTOCOL = 2
                with pd.HDFStore(storepath, mode='w') as store:
                    store[self.DF_NAME] = df
            finally:
                cPickle.HIGHEST_PROTOCOL = old_protocol

{'fraction_non_alphanumeric': 0.0769992795389049, 'fraction_numerical': 0.004592939481268011, 'mean_word_length': 2.2104654524429024, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3912163', 'n_tokens_mistral': 3395, 'n_tokens_neox': 3246, 'n_words': 907}