<commit_before>from __future__ import absolute_import

try:
    import itertools.izip as zip
except ImportError:
    pass

import numpy as np
import pandas as pd
import dask.dataframe as dd
from dask.dataframe import DataFrame
from dask.dataframe.core import Scalar

from .. import util
from ..element import Element
from ..ndmapping import NdMapping, item_check
from .interface import Interface
from .pandas import PandasInterface


class DaskInterface(PandasInterface):
    """
    The DaskInterface allows a Dataset objects to wrap a dask
    DataFrame object. Using dask allows loading data lazily
    and performing out-of-core operations on the data, making
    it possible to work on datasets larger than memory.

    The DaskInterface covers almost the complete API exposed
    by the PandasInterface with two notable exceptions:

    1) Sorting is not supported and any attempt at sorting will
       be ignored with an warning.
    2) Dask does not easily support adding a new column to an existing
       dataframe unless it is a scalar, add_dimension will therefore
       error when supplied a non-scalar value.
    4) Not all functions can be easily applied to a dask dataframe so
       some functions applied with aggregate and reduce will not work.
    """

    types = (DataFrame,)

    datatype = 'dask'

    default_partitions = 100

    @classmethod
    def init(cls, eltype, data, kdims, vdims):
        data, kdims, vdims = PandasInterface.init(eltype, data, kdims, vdims)
        if not isinstance(data, DataFrame):
            data = dd.from_pandas(data, npartitions=cls.default_partitions, sort=False)
        return data, kdims, vdims

    @classmethod
    def shape(cls, dataset):
        return (len(dataset.data), len(dataset.data.columns))

    @classmethod
    def range(cls, columns, dimension):
        column = columns.data[columns.get_dimension(dimension).name]
        if column.dtype.kind == 'O':
            column = np.sort(column[column.notnull()].compute())
            return column[0], column[-1]
        else:
            return dd.compute(column.min(), column.max())

    @classmethod
    def sort(cls, columns, by=[]):
        columns.warning('Dask dataframes do not support sorting')
        return columns.data

    @classmethod
    def values(cls, columns, dim, expanded=True, flat=True):
        data = columns.data[dim]
        if not expanded:
            data = data.unique()
        return data.compute().values

    @classmethod
    def select_mask(cls, dataset, selection):
        """
        Given a Dataset object and a dictionary with dimension keys and
        selection keys (i.e tuple ranges, slices, sets, lists or literals)
        return a boolean mask over the rows in the Dataset object that
        have been selected.
        """
        select_mask = None
        for dim, k in selection.items():
            if isinstance(k, tuple):
                k = slice(*k)
            masks = []
            series = dataset.data[dim]
            if isinstance(k, slice):
                if k.start is not None:
                    masks.append(k.start <= series)
                if k.stop is not None:
                    masks.append(series < k.stop)
            elif isinstance(k, (set, list)):
                iter_slc = None
                for ik in k:
                    mask = series == ik
                    if iter_slc is None:
                        iter_slc = mask
                    else:
                        iter_slc |= mask
                masks.append(iter_slc)
            elif callable(k):
                masks.append(k(series))
            else:
                masks.append(series == k)
            for mask in masks:
                if select_mask:
                    select_mask &= mask
                else:
                    select_mask = mask
        return select_mask

    @classmethod
    def select(cls, columns, selection_mask=None, **selection):
        df = columns.data
        if selection_mask is not None:
            return df[selection_mask]
        selection_mask = cls.select_mask(columns, selection)
        indexed = cls.indexed(columns, selection)
        df = df if selection_mask is None else df[selection_mask]
        if indexed and len(df) == 1:
            return df[columns.vdims[0].name].compute().iloc[0]
        return df
    
    @classmethod
    def groupby(cls, columns, dimensions, container_type, group_type, **kwargs):
        index_dims = [columns.get_dimension(d) for d in dimensions]
        element_dims = [kdim for kdim in columns.kdims
                        if kdim not in index_dims]

        group_kwargs = {}
        if group_type != 'raw' and issubclass(group_type, Element):
            group_kwargs = dict(util.get_param_values(columns),
                                kdims=element_dims)
        group_kwargs.update(kwargs)

        data = []
        groupby = columns.data.groupby(dimensions)
        ind_array = columns.data[dimensions].compute().values
        indices = (tuple(ind) for ind in ind_array)
        for coord in util.unique_iterator(indices):
            if any(isinstance(c, float) and np.isnan(c) for c in coord):
                continue
            if len(coord) == 1:
                coord = coord[0]
            group = group_type(groupby.get_group(coord), **group_kwargs)
            data.append((coord, group))
        if issubclass(container_type, NdMapping):<commit_msg>Optimized dask interface groupby implementation<commit_after>from __future__ import absolute_import

try:
    import itertools.izip as zip
except ImportError:
    pass

import numpy as np
import pandas as pd
import dask.dataframe as dd
from dask.dataframe import DataFrame
from dask.dataframe.core import Scalar

from .. import util
from ..element import Element
from ..ndmapping import NdMapping, item_check
from .interface import Interface
from .pandas import PandasInterface


class DaskInterface(PandasInterface):
    """
    The DaskInterface allows a Dataset objects to wrap a dask
    DataFrame object. Using dask allows loading data lazily
    and performing out-of-core operations on the data, making
    it possible to work on datasets larger than memory.

    The DaskInterface covers almost the complete API exposed
    by the PandasInterface with two notable exceptions:

    1) Sorting is not supported and any attempt at sorting will
       be ignored with an warning.
    2) Dask does not easily support adding a new column to an existing
       dataframe unless it is a scalar, add_dimension will therefore
       error when supplied a non-scalar value.
    4) Not all functions can be easily applied to a dask dataframe so
       some functions applied with aggregate and reduce will not work.
    """

    types = (DataFrame,)

    datatype = 'dask'

    default_partitions = 100

    @classmethod
    def init(cls, eltype, data, kdims, vdims):
        data, kdims, vdims = PandasInterface.init(eltype, data, kdims, vdims)
        if not isinstance(data, DataFrame):
            data = dd.from_pandas(data, npartitions=cls.default_partitions, sort=False)
        return data, kdims, vdims

    @classmethod
    def shape(cls, dataset):
        return (len(dataset.data), len(dataset.data.columns))

    @classmethod
    def range(cls, columns, dimension):
        column = columns.data[columns.get_dimension(dimension).name]
        if column.dtype.kind == 'O':
            column = np.sort(column[column.notnull()].compute())
            return column[0], column[-1]
        else:
            return dd.compute(column.min(), column.max())

    @classmethod
    def sort(cls, columns, by=[]):
        columns.warning('Dask dataframes do not support sorting')
        return columns.data

    @classmethod
    def values(cls, columns, dim, expanded=True, flat=True):
        data = columns.data[dim]
        if not expanded:
            data = data.unique()
        return data.compute().values

    @classmethod
    def select_mask(cls, dataset, selection):
        """
        Given a Dataset object and a dictionary with dimension keys and
        selection keys (i.e tuple ranges, slices, sets, lists or literals)
        return a boolean mask over the rows in the Dataset object that
        have been selected.
        """
        select_mask = None
        for dim, k in selection.items():
            if isinstance(k, tuple):
                k = slice(*k)
            masks = []
            series = dataset.data[dim]
            if isinstance(k, slice):
                if k.start is not None:
                    masks.append(k.start <= series)
                if k.stop is not None:
                    masks.append(series < k.stop)
            elif isinstance(k, (set, list)):
                iter_slc = None
                for ik in k:
                    mask = series == ik
                    if iter_slc is None:
                        iter_slc = mask
                    else:
                        iter_slc |= mask
                masks.append(iter_slc)
            elif callable(k):
                masks.append(k(series))
            else:
                masks.append(series == k)
            for mask in masks:
                if select_mask:
                    select_mask &= mask
                else:
                    select_mask = mask
        return select_mask

    @classmethod
    def select(cls, columns, selection_mask=None, **selection):
        df = columns.data
        if selection_mask is not None:
            return df[selection_mask]
        selection_mask = cls.select_mask(columns, selection)
        indexed = cls.indexed(columns, selection)
        df = df if selection_mask is None else df[selection_mask]
        if indexed and len(df) == 1:
            return df[columns.vdims[0].name].compute().iloc[0]
        return df
    
    @classmethod
    def groupby(cls, columns, dimensions, container_type, group_type, **kwargs):
        index_dims = [columns.get_dimension(d) for d in dimensions]
        element_dims = [kdim for kdim in columns.kdims
                        if kdim not in index_dims]

        group_kwargs = {}
        if group_type != 'raw' and issubclass(group_type, Element):
            group_kwargs = dict(util.get_param_values(columns),
                                kdims=element_dims)
        group_kwargs.update(kwargs)

        data = []
        groupby = columns.data.groupby(dimensions)
        if len(dimensions) == 1:
            column = columns.data[dimensions[0]]
            if column.dtype.name == 'category':
                indices = ((ind,) for ind in column.cat.categories)
            else:
                indices = ((ind,) for ind in column.unique().compute())
        else:
            group_tuples = columns.data[dimensions].itertuple()
            indices = util.unique_iterator(ind[1:] for ind in group_tuples)
        for coord in indices:
            if any(isinstance(c, float) and np.isnan(c) for c in coord):
                continue
            if len(coord) == 1:
                coord = coord[0]
            group = group_type(groupby.get_group(coord), **group_kwargs)
            data.append((coord, group))
        if issubclass(container_type, NdMapping):
{'fraction_non_alphanumeric': 0.08047000178030977, 'fraction_numerical': 0.0025814491721559553, 'mean_word_length': 2.046366594360087, 'pattern_counts': {'":': 0, '<': 7, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4596681', 'n_tokens_mistral': 3061, 'n_tokens_neox': 2991, 'n_words': 995}