<commit_before>import uuid
from time import gmtime, strftime

from celery.task import task
from pandas import rolling_window

from bamboo.core.calculator import Calculator
from bamboo.core.frame import BambooFrame, BAMBOO_RESERVED_KEY_PREFIX,\
    DATASET_ID, INDEX, PARENT_DATASET_ID
from bamboo.core.summary import summarize
from bamboo.lib.async import call_async
from bamboo.lib.exceptions import ArgumentError
from bamboo.lib.io import ImportableDataset
from bamboo.lib.query_args import QueryArgs
from bamboo.lib.schema_builder import Schema
from bamboo.lib.utils import to_list
from bamboo.models.abstract_model import AbstractModel
from bamboo.models.calculation import Calculation
from bamboo.models.observation import Observation


@task(ignore_result=True)
def delete_task(dataset):
    """Background task to delete dataset and its associated observations."""
    Observation.delete_all(dataset)
    super(dataset.__class__, dataset).delete({DATASET_ID: dataset.dataset_id})
    Observation.delete_encoding(dataset)


class Dataset(AbstractModel, ImportableDataset):

    __collectionname__ = 'datasets'

    # caching keys
    STATS = '_stats'
    ALL = '_all'

    # metadata
    AGGREGATED_DATASETS = BAMBOO_RESERVED_KEY_PREFIX + 'linked_datasets'
    ATTRIBUTION = 'attribution'
    CREATED_AT = 'created_at'
    DESCRIPTION = 'description'
    ID = 'id'
    JOINED_DATASETS = 'joined_datasets'
    LABEL = 'label'
    LICENSE = 'license'
    NUM_COLUMNS = 'num_columns'
    NUM_ROWS = 'num_rows'
    MERGED_DATASETS = 'merged_datasets'
    PARENT_IDS = 'parent_ids'
    PENDING_UPDATES = 'pending_updates'
    SCHEMA = 'schema'
    UPDATED_AT = 'updated_at'

    def __init__(self, record=None):
        super(Dataset, self).__init__(record)
        self.__dframe = None

    @property
    def dataset_id(self):
        return self.record[DATASET_ID]

    @property
    def num_columns(self):
        return self.record.get(self.NUM_COLUMNS, 0)

    @property
    def num_rows(self):
        return self.record.get(self.NUM_ROWS, 0)

    @property
    def schema(self):
        schema_dict = {}

        if self.record:
            schema_dict = self.record.get(self.SCHEMA)

        return Schema.safe_init(schema_dict)

    @property
    def columns(self):
        return self.schema.keys() if self.num_rows else []

    @property
    def labels(self):
        return [column[self.LABEL] for column in self.schema.values()]

    @property
    def label(self):
        return self.record.get(self.LABEL)

    @property
    def description(self):
        return self.record.get(self.DESCRIPTION)

    @property
    def license(self):
        return self.record.get(self.LICENSE)

    @property
    def attribution(self):
        return self.record.get(self.ATTRIBUTION)

    @property
    def stats(self):
        return self.record.get(self.STATS, {})

    @property
    def aggregated_datasets_dict(self):
        return self.record.get(self.AGGREGATED_DATASETS, {})

    @property
    def aggregated_datasets(self):
        return [(self.split_groups(group), self.find_one(_id)) for (
            group, _id) in self.aggregated_datasets_dict.items()]

    @property
    def joined_dataset_ids(self):
        return [
            tuple(_list) for _list in self.record.get(self.JOINED_DATASETS, [])
        ]

    @property
    def joined_datasets(self):
        # TODO: fetch all datasets in single DB call
        # (let Dataset.find take a list of IDs)
        return [
            (direction, self.find_one(other_dataset_id), on,
             self.find_one(joined_dataset_id))
            for direction, other_dataset_id, on, joined_dataset_id in
            self.joined_dataset_ids]

    @property
    def merged_dataset_info(self):
        return self.record.get(self.MERGED_DATASETS, [])

    @property
    def merged_dataset_ids(self):
        results = self.merged_dataset_info
        return zip(*results)[-1] if results else results

    @property
    def merged_datasets(self):
        return self.__linked_datasets(self.merged_dataset_ids)

    @property
    def merged_datasets_with_map(self):
        results = self.merged_dataset_info

        if len(results):
            mappings, ids = zip(*results)
            results = zip(mappings, self.__linked_datasets(ids))

        return results

    @property
    def pending_updates(self):
        return self.record[self.PENDING_UPDATES]

    @property
    def updatable_keys(self):
        return [self.LABEL, self.DESCRIPTION, self.LICENSE, self.ATTRIBUTION]

    @property
    def on_columns_for_rhs_of_joins(self):
        return [on for direction, _, on, __ in
                self.joined_datasets if direction == 'left']

    @property
    def __is_cached(self):
        return self.__dframe is not None

    def __linked_datasets(self, ids):
        return [self.find_one(_id) for _id in ids]

    def is_factor(self, col):
        return self.schema.is_dimension(col)
<commit_msg>allow date columns in multigroup<commit_after>import uuid
from time import gmtime, strftime

from celery.task import task
from pandas import rolling_window

from bamboo.core.calculator import Calculator
from bamboo.core.frame import BambooFrame, BAMBOO_RESERVED_KEY_PREFIX,\
    DATASET_ID, INDEX, PARENT_DATASET_ID
from bamboo.core.summary import summarize
from bamboo.lib.async import call_async
from bamboo.lib.exceptions import ArgumentError
from bamboo.lib.io import ImportableDataset
from bamboo.lib.query_args import QueryArgs
from bamboo.lib.schema_builder import Schema
from bamboo.lib.utils import to_list
from bamboo.models.abstract_model import AbstractModel
from bamboo.models.calculation import Calculation
from bamboo.models.observation import Observation


@task(ignore_result=True)
def delete_task(dataset):
    """Background task to delete dataset and its associated observations."""
    Observation.delete_all(dataset)
    super(dataset.__class__, dataset).delete({DATASET_ID: dataset.dataset_id})
    Observation.delete_encoding(dataset)


class Dataset(AbstractModel, ImportableDataset):

    __collectionname__ = 'datasets'

    # caching keys
    STATS = '_stats'
    ALL = '_all'

    # metadata
    AGGREGATED_DATASETS = BAMBOO_RESERVED_KEY_PREFIX + 'linked_datasets'
    ATTRIBUTION = 'attribution'
    CREATED_AT = 'created_at'
    DESCRIPTION = 'description'
    ID = 'id'
    JOINED_DATASETS = 'joined_datasets'
    LABEL = 'label'
    LICENSE = 'license'
    NUM_COLUMNS = 'num_columns'
    NUM_ROWS = 'num_rows'
    MERGED_DATASETS = 'merged_datasets'
    PARENT_IDS = 'parent_ids'
    PENDING_UPDATES = 'pending_updates'
    SCHEMA = 'schema'
    UPDATED_AT = 'updated_at'

    def __init__(self, record=None):
        super(Dataset, self).__init__(record)
        self.__dframe = None

    @property
    def dataset_id(self):
        return self.record[DATASET_ID]

    @property
    def num_columns(self):
        return self.record.get(self.NUM_COLUMNS, 0)

    @property
    def num_rows(self):
        return self.record.get(self.NUM_ROWS, 0)

    @property
    def schema(self):
        schema_dict = {}

        if self.record:
            schema_dict = self.record.get(self.SCHEMA)

        return Schema.safe_init(schema_dict)

    @property
    def columns(self):
        return self.schema.keys() if self.num_rows else []

    @property
    def labels(self):
        return [column[self.LABEL] for column in self.schema.values()]

    @property
    def label(self):
        return self.record.get(self.LABEL)

    @property
    def description(self):
        return self.record.get(self.DESCRIPTION)

    @property
    def license(self):
        return self.record.get(self.LICENSE)

    @property
    def attribution(self):
        return self.record.get(self.ATTRIBUTION)

    @property
    def stats(self):
        return self.record.get(self.STATS, {})

    @property
    def aggregated_datasets_dict(self):
        return self.record.get(self.AGGREGATED_DATASETS, {})

    @property
    def aggregated_datasets(self):
        return [(self.split_groups(group), self.find_one(_id)) for (
            group, _id) in self.aggregated_datasets_dict.items()]

    @property
    def joined_dataset_ids(self):
        return [
            tuple(_list) for _list in self.record.get(self.JOINED_DATASETS, [])
        ]

    @property
    def joined_datasets(self):
        # TODO: fetch all datasets in single DB call
        # (let Dataset.find take a list of IDs)
        return [
            (direction, self.find_one(other_dataset_id), on,
             self.find_one(joined_dataset_id))
            for direction, other_dataset_id, on, joined_dataset_id in
            self.joined_dataset_ids]

    @property
    def merged_dataset_info(self):
        return self.record.get(self.MERGED_DATASETS, [])

    @property
    def merged_dataset_ids(self):
        results = self.merged_dataset_info
        return zip(*results)[-1] if results else results

    @property
    def merged_datasets(self):
        return self.__linked_datasets(self.merged_dataset_ids)

    @property
    def merged_datasets_with_map(self):
        results = self.merged_dataset_info

        if len(results):
            mappings, ids = zip(*results)
            results = zip(mappings, self.__linked_datasets(ids))

        return results

    @property
    def pending_updates(self):
        return self.record[self.PENDING_UPDATES]

    @property
    def updatable_keys(self):
        return [self.LABEL, self.DESCRIPTION, self.LICENSE, self.ATTRIBUTION]

    @property
    def on_columns_for_rhs_of_joins(self):
        return [on for direction, _, on, __ in
                self.joined_datasets if direction == 'left']

    @property
    def __is_cached(self):
        return self.__dframe is not None

    def __linked_datasets(self, ids):
        return [self.find_one(_id) for _id in ids]

    def is_dimension(self, col):
        return self.schema.is_dimension(col)

    def is_factor(self, col):
        return self.is_dimension(col) or self.schema.is_date_simpletype(col)

{'fraction_non_alphanumeric': 0.0890764647467726, 'fraction_numerical': 0.0005958291956305859, 'mean_word_length': 3.0608870967741937, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7523279', 'n_tokens_mistral': 3214, 'n_tokens_neox': 3077, 'n_words': 695}