<commit_before># internal imports
import ctable_ext

# external imports
import numpy as np
import bcolz
from collections import namedtuple
import os
from bquery.ctable_ext import \
    SUM, COUNT, COUNT_NA, COUNT_DISTINCT, SORTED_COUNT_DISTINCT

class ctable(bcolz.ctable):
    def cache_factor(self, col_list, refresh=False):
        """
        Existing todos here are: these should be hidden helper carrays
        As in: not normal columns that you would normally see as a user

        The factor (label index) carray is as long as the original carray
        (and the rest of the table therefore)
        But the (unique) values carray is not as long (as long as the number
        of unique values)

        :param col_list:
        :param refresh:
        :return:
        """

        if not self.rootdir:
            raise TypeError('Only out-of-core ctables can have '
                            'factorization caching at the moment')

        for col in col_list:

            col_rootdir = self[col].rootdir
            col_factor_rootdir = col_rootdir + '.factor'
            col_values_rootdir = col_rootdir + '.values'

            # create cache if needed
            if refresh or not os.path.exists(col_factor_rootdir):
                carray_factor = \
                    bcolz.carray([], dtype='int64', expectedlen=self.size,
                                   rootdir=col_factor_rootdir, mode='w')
                _, values = \
                    ctable_ext.factorize(self[col], labels=carray_factor)
                carray_factor.flush()
                carray_values = \
                    bcolz.carray(values.values(), dtype=self[col].dtype,
                                 rootdir=col_values_rootdir, mode='w')
                carray_values.flush()

    def groupby(self, groupby_cols, agg_list, bool_arr=None, rootdir=None,
                agg_method='sum'):
        """
        Aggregate the ctable

        groupby_cols: a list of columns to groupby over
        agg_list: the aggregation operations, which can be:
         - a straight forward sum of a list columns with a
           similarly named output: ['m1', 'm2', ...]
         - a list of new column input/output settings
           [['mnew1', 'm1'], ['mnew2', 'm2], ...]
         - a list that includes the type of aggregation for each column, i.e.
           [['mnew1', 'm1', 'sum'], ['mnew2', 'm1, 'avg'], ...]

        Currently supported aggregation operations are:
        - sum
        - sum_na (that checks for nan values and excludes them)
        - To be added: mean, mean_na (and perhaps standard deviation etc)

        boolarr: to be added (filtering the groupby factorization input)
        rootdir: the aggregation ctable rootdir

        """
        map_agg_method = {
            'sum': SUM,
            'count': COUNT,
            'count_na': COUNT_NA,
            'count_distinct': COUNT_DISTINCT,
            'sorted_count_distinct': SORTED_COUNT_DISTINCT,
        }
        _agg_method = map_agg_method[agg_method]

        if not agg_list:
            raise AttributeError('One or more aggregation operations '
                                 'need to be defined')

        factor_list, values_list = self.factorize_groupby_cols(groupby_cols)

        factor_carray, nr_groups, skip_key = \
            self.make_group_index(factor_list, values_list, groupby_cols,
                                  len(self), bool_arr)

        ct_agg, dtype_list, agg_ops = \
            self.create_agg_ctable(groupby_cols, agg_list, nr_groups, rootdir)

        # perform aggregation
        ctable_ext.aggregate_groups_by_iter_2(self, ct_agg, nr_groups, skip_key,
                                              factor_carray, groupby_cols,
                                              agg_ops, dtype_list,<commit_msg>Added unique (distinct) functionality that uses factorization (and if available, caching)<commit_after># internal imports
import ctable_ext

# external imports
import numpy as np
import bcolz
from collections import namedtuple
import os
from bquery.ctable_ext import \
    SUM, COUNT, COUNT_NA, COUNT_DISTINCT, SORTED_COUNT_DISTINCT

class ctable(bcolz.ctable):
    def cache_factor(self, col_list, refresh=False):
        """
        Existing todos here are: these should be hidden helper carrays
        As in: not normal columns that you would normally see as a user

        The factor (label index) carray is as long as the original carray
        (and the rest of the table therefore)
        But the (unique) values carray is not as long (as long as the number
        of unique values)

        :param col_list:
        :param refresh:
        :return:
        """

        if not self.rootdir:
            raise TypeError('Only out-of-core ctables can have '
                            'factorization caching at the moment')

        for col in col_list:

            col_rootdir = self[col].rootdir
            col_factor_rootdir = col_rootdir + '.factor'
            col_values_rootdir = col_rootdir + '.values'

            # create cache if needed
            if refresh or not os.path.exists(col_factor_rootdir):
                carray_factor = \
                    bcolz.carray([], dtype='int64', expectedlen=self.size,
                                   rootdir=col_factor_rootdir, mode='w')
                _, values = \
                    ctable_ext.factorize(self[col], labels=carray_factor)
                carray_factor.flush()
                carray_values = \
                    bcolz.carray(values.values(), dtype=self[col].dtype,
                                 rootdir=col_values_rootdir, mode='w')
                carray_values.flush()

    def unique(self, col_or_col_list):
        """
        Return a list of unique values of a column or a list of lists of column list

        :param col_or_col_list: a column or a list of columns
        :return:
        """

        if isinstance(col_or_col_list, list):
            col_is_list = True
            col_list = col_or_col_list
        else:
            col_is_list = False
            col_list = [col_or_col_list]

        output = []

        for col in col_list:

            if self.rootdir:
                col_values_rootdir = self[col].rootdir + '.values'
            else:
                col_values_rootdir = None

            if col_values_rootdir and os.path.exists(col_values_rootdir):
                # retrieve values from existing disk-based factorization
                carray_values = bcolz.carray(rootdir=col_values_rootdir)
                values = list(carray_values)
            else:
                # factorize on-the-fly
                _, values = ctable_ext.factorize(self[col])
                values = values.values()

            output.append(values)

        if not col_is_list:
            output = output[0]

        return output

    def groupby(self, groupby_cols, agg_list, bool_arr=None, rootdir=None,
                agg_method='sum'):
        """

        Aggregate the ctable

        groupby_cols: a list of columns to groupby over
        agg_list: the aggregation operations, which can be:
         - a straight forward sum of a list columns with a
           similarly named output: ['m1', 'm2', ...]
         - a list of new column input/output settings
           [['mnew1', 'm1'], ['mnew2', 'm2], ...]
         - a list that includes the type of aggregation for each column, i.e.
           [['mnew1', 'm1', 'sum'], ['mnew2', 'm1, 'avg'], ...]

        Currently supported aggregation operations are:
        - sum
        - sum_na (that checks for nan values and excludes them)
        - To be added: mean, mean_na (and perhaps standard deviation etc)

        boolarr: to be added (filtering the groupby factorization input)
        rootdir: the aggregation ctable rootdir

        """
        # TODO: change aggregation types to method as described in "a list with the type of aggregation for each column"
        map_agg_method = {
            'sum': SUM,
            'count': COUNT,
            'count_na': COUNT_NA,
            'count_distinct': COUNT_DISTINCT,
            'sorted_count_distinct': SORTED_COUNT_DISTINCT,
        }
        _agg_method = map_agg_method[agg_method]

        if not agg_list:
            raise AttributeError('One or more aggregation operations '
                                 'need to be defined')

        factor_list, values_list = self.factorize_groupby_cols(groupby_cols)

        factor_carray, nr_groups, skip_key = \
            self.make_group_index(factor_list, values_list, groupby_cols,
                                  len(self), bool_arr)

        ct_agg, dtype_list, agg_ops = \
            self.create_agg_ctable(groupby_cols, agg_list, nr_groups, rootdir)

        # perform aggregation
        ctable_ext.aggregate_groups_by_iter_2(self, ct_agg, nr_groups, skip_key,
                                              factor_carray, groupby_cols,
                                              agg_ops, dtype_list,
{'fraction_non_alphanumeric': 0.07680372381691233, 'fraction_numerical': 0.0029923528759835975, 'mean_word_length': 1.8547927870926921, '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': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4391914', 'n_tokens_mistral': 2572, 'n_tokens_neox': 2430, 'n_words': 732}