<commit_before>from itertools import izip, imap, chain, tee
import operator
import collections
import types
import numpy as np
import pandas as pd
import psycopg2
from ml.utils.seq import grouper_chunk
from ml.utils.numeric_functions import max_type


def choice(operator):
    def inner(fn):
        def view(x, y):
            if isinstance(x, IterLayer) and isinstance(y, IterLayer):
                return x.stream_operation(operator, y)
            elif isinstance(y, collections.Iterable):
                return x.stream_operation(operator, IterLayer(y, length=x.length))
            else:
                return x.scalar_operation(operator, y)
        return view
    return inner


class IterLayer(object):
    def __init__(self, fn_iter, shape=None, dtype=None, type_elem=None,
                has_chunks=False, chunks_size=0, length=None):
        dtypes = None
        if isinstance(fn_iter, types.GeneratorType) or isinstance(fn_iter, psycopg2.extensions.cursor):
            _fn_iter = fn_iter
        elif isinstance(fn_iter, IterLayer):
            _fn_iter = fn_iter
        elif isinstance(fn_iter, pd.DataFrame):
            _fn_iter = (e for e in fn_iter.values)
            dtypes = zip(fn_iter.columns.values, fn_iter.dtypes.values)
        else:
            _fn_iter = (e for e in fn_iter)

        self.it = _fn_iter
        self.pushedback = []
        self.chunks_size = chunks_size
        self.has_chunks = has_chunks
        if shape is not None:
            self.length = shape[0]
        else:
            self.length = length

        if dtype is None:
            #obtain dtype, shape, global_dtype and type_elem
            self.chunk_taste(dtypes)
        else:
            if hasattr(dtype, '__iter__'):
                self.global_dtype = self._get_global_dtype(dtype)
            else:
                self.global_dtype = dtype
            self.dtype = dtype
            self.shape = shape
            self.type_elem = np.ndarray if not isinstance(self.dtype, list) else pd.DataFrame
    
    def pushback(self, val):
        self.pushedback.append(val)

    def chunk_taste(self, dtypes):
        """Check for the dtype and global dtype in a chunk"""
        chunk = next(self)
        if dtypes is not None:
            self.dtype = dtypes
            global_dtype = self._get_global_dtype(dtypes)
        elif isinstance(chunk, pd.DataFrame):
            self.dtype = []
            for c, cdtype in zip(chunk.columns.values, chunk.dtypes.values):
                self.dtype.append((c, cdtype))
            global_dtype = self._get_global_dtype(self.dtype)
        elif isinstance(chunk, np.ndarray):
            self.dtype = chunk.dtype
            global_dtype = self.dtype
        else:#scalars
            if type(chunk).__module__ == '__builtin__':
                if hasattr(chunk, '__iter__'):
                    type_e = max_type(chunk)
                    if isinstance(type_e, list) or isinstance(type_e, tuple) or\
                        type_e == str or type_e == unicode:
                        self.dtype = "|O"
                    else:
                        self.dtype = type_e
                else:
                    self.dtype = type(chunk)
            else:
                self.dtype = chunk.dtype
            global_dtype = self.dtype
        
        try:
            shape = chunk.shape
        except AttributeError:
            if hasattr(chunk, '__iter__'):
                shape = (len(chunk),)
            else:
                shape = (1,)

        if self.has_chunks:
            shape = list(shape[1:])
            if len(shape) == 0:
                shape = [0]
            self.shape = [self.length] + list(shape)
        else:
            if len(shape) > 1:
                self.shape = [self.length] + list(shape[1:])
            else:
                self.shape = [self.length] + list(shape)

        self.global_dtype = global_dtype
        self.pushback(chunk)
        self.type_elem = type(chunk)

    def _get_global_dtype(self, dtype):
        global_dtype = None
        for c, cdtype in dtype:
            if cdtype == np.dtype("|O"):
                global_dtype = cdtype
                break
            else:
                global_dtype = cdtype
        return global_dtype

    def to_chunks(self, chunks_size, dtype=None):
        if self.has_chunks is False:
            if dtype is None:
                dtype = self.dtype
            return IterLayer(self.gen_chunks(chunks_size, dtype), shape=self.shape, 
                chunks_size=chunks_size, has_chunks=True, dtype=dtype)
        else:
            return self
    
    def gen_chunks(self, chunks_size, dtype):
        if chunks_size < 1:
            chunks_size = self.shape[0]

        if len(self.shape) == 1:
            chunk_shape = [chunks_size]
        else:
            i_features = self.shape[1:]
            chunk_shape = [chunks_size] + list(i_features)

        if not isinstance(dtype, list):
            for smx in grouper_chunk(chunks_size, self):
                smx_a = np.empty(chunk_shape, dtype=np.dtype(dtype))
                for i, row in enumerate(smx):
                    try:
                        smx_a[i] = row
                    except ValueError:
                        smx_a[i] = row[0]
                if i + 1 < chunks_size:
                    yield smx_a[:i+1]
                else:
                    yield smx_a
        else:
            for smx in grouper_chunk(chunks_size, self):
                smx_a = pd.DataFrame(index=np.arange(0, chunk_shape[0]), 
                    dtype=dtype[0][1], columns=[c for c, _ in dtype])
                for i, row in enumerate(smx):
                    #try:
                    smx_a.values[i] = row
                    #except ValueError:
                    #    smx_a.values[i] = row[0]
                if i + 1 < chunks_size:
                    yield smx_a.iloc[:i+1]
                else:
                    yield smx_a
                
    def flat(self):
        def _iter():
            if self.type_elem == np.ndarray:
                for chunk in self:
                    for e in chunk.reshape(-1):
                        yield e
            elif self.type_elem == pd.DataFrame:
                for chunk in self:
                    for e in chunk.values.reshape(-1):
                        yield e
            elif self.type_elem.__module__ == '__builtin__':
                for e in chain.from_iterable(self):
                    yield e
        
        shape = (reduce(operator.mul, self.shape),)
        it = IterLayer(_iter(), shape=shape, dtype=self.dtype, length=self.length)
        if self.has_chunks:            
            return it.to_chunks(self.chunks_size, dtype=self.global_dtype)
        else:<commit_msg>fixed performance for chunks builder and dtype selection<commit_after>from itertools import izip, imap, chain, tee
import operator
import collections
import types
import numpy as np
import pandas as pd
import psycopg2
from ml.utils.seq import grouper_chunk
from ml.utils.numeric_functions import max_type


def choice(operator):
    def inner(fn):
        def view(x, y):
            if isinstance(x, IterLayer) and isinstance(y, IterLayer):
                return x.stream_operation(operator, y)
            elif isinstance(y, collections.Iterable):
                return x.stream_operation(operator, IterLayer(y, length=x.length))
            else:
                return x.scalar_operation(operator, y)
        return view
    return inner


class IterLayer(object):
    def __init__(self, fn_iter, shape=None, dtype=None, type_elem=None,
                has_chunks=False, chunks_size=0, length=None):
        dtypes = None
        if isinstance(fn_iter, types.GeneratorType) or isinstance(fn_iter, psycopg2.extensions.cursor):
            _fn_iter = fn_iter
        elif isinstance(fn_iter, IterLayer):
            _fn_iter = fn_iter
        elif isinstance(fn_iter, pd.DataFrame):
            _fn_iter = (e for e in fn_iter.values)
            dtypes = zip(fn_iter.columns.values, fn_iter.dtypes.values)
        else:
            _fn_iter = (e for e in fn_iter)

        self.it = _fn_iter
        self.pushedback = []
        self.chunks_size = chunks_size
        self.has_chunks = has_chunks
        if shape is not None:
            self.length = shape[0]
        else:
            self.length = length

        if dtype is None:
            #obtain dtype, shape, global_dtype and type_elem
            self.chunk_taste(dtypes)
        else:
            if hasattr(dtype, '__iter__'):
                self.global_dtype = self._get_global_dtype(dtype)
            else:
                self.global_dtype = dtype
            self.dtype = dtype
            self.shape = shape
            self.type_elem = np.ndarray if not isinstance(self.dtype, list) else pd.DataFrame
    
    def pushback(self, val):
        self.pushedback.append(val)

    def chunk_taste(self, dtypes):
        """Check for the dtype and global dtype in a chunk"""
        chunk = next(self)
        if dtypes is not None:
            self.dtype = dtypes
            global_dtype = self._get_global_dtype(dtypes)
        elif isinstance(chunk, pd.DataFrame):
            self.dtype = []
            for c, cdtype in zip(chunk.columns.values, chunk.dtypes.values):
                self.dtype.append((c, cdtype))
            global_dtype = self._get_global_dtype(self.dtype)
        elif isinstance(chunk, np.ndarray):
            self.dtype = chunk.dtype
            global_dtype = self.dtype
        else:#scalars
            if type(chunk).__module__ == '__builtin__':
                if hasattr(chunk, '__iter__'):
                    type_e = max_type(chunk)
                    if isinstance(type_e, list) or isinstance(type_e, tuple) or\
                        type_e == str or type_e == unicode:
                        self.dtype = "|O"
                    else:
                        self.dtype = type_e
                else:
                    self.dtype = type(chunk)
            else:
                self.dtype = chunk.dtype
            global_dtype = self.dtype
        
        try:
            shape = chunk.shape
        except AttributeError:
            if hasattr(chunk, '__iter__'):
                shape = (len(chunk),)
            else:
                shape = (1,)

        if self.has_chunks:
            shape = list(shape[1:])
            if len(shape) == 0:
                shape = [0]
            self.shape = [self.length] + list(shape)
        else:
            if len(shape) > 1:
                self.shape = [self.length] + list(shape[1:])
            else:
                self.shape = [self.length] + list(shape)

        self.global_dtype = global_dtype
        self.pushback(chunk)
        self.type_elem = type(chunk)

    def _get_global_dtype(self, dtype):
        sizeof = [(np.dtype(type(cdtype)), cdtype) for _, cdtype in dtype]
        return max(sizeof, key=lambda x: x[0])[1]

    def to_chunks(self, chunks_size, dtype=None):
        if self.has_chunks is False:
            if dtype is None:
                dtype = self.dtype
            return IterLayer(self.gen_chunks(chunks_size, dtype), shape=self.shape, 
                chunks_size=chunks_size, has_chunks=True, dtype=dtype)
        else:
            return self
    
    def gen_chunks(self, chunks_size, dtype):
        if chunks_size < 1:
            chunks_size = self.shape[0]

        if len(self.shape) == 1:
            chunk_shape = [chunks_size]
        else:
            i_features = self.shape[1:]
            chunk_shape = [chunks_size] + list(i_features)

        if not isinstance(dtype, list):
            for smx in grouper_chunk(chunks_size, self):
                smx_a = np.empty(chunk_shape, dtype=np.dtype(dtype))
                for i, row in enumerate(smx):
                    try:
                        smx_a[i] = row
                    except ValueError:
                        smx_a[i] = row[0]
                if i + 1 < chunks_size:
                    yield smx_a[:i+1]
                else:
                    yield smx_a
        else:
            cdtype = dict(dtype)
            columns = [c for c, _ in dtype]
            for smx in grouper_chunk(chunks_size, self):
                x = np.empty(chunk_shape[0], dtype=dtype)
                for i, row in enumerate(smx):
                    try:
                        x[i] = row
                    except ValueError:
                        x[i] = row[0]
                smx_a = pd.DataFrame(x,
                    index=np.arange(0, chunk_shape[0]), 
                    columns=columns)
                
                if i + 1 < chunks_size:
                    yield smx_a.iloc[:i+1]
                else:
                    yield smx_a
                
    def flat(self):
        def _iter():
            if self.type_elem == np.ndarray:
                for chunk in self:
                    for e in chunk.reshape(-1):
                        yield e
            elif self.type_elem == pd.DataFrame:
                for chunk in self:
                    for e in chunk.values.reshape(-1):
                        yield e
            elif self.type_elem.__module__ == '__builtin__':
                for e in chain.from_iterable(self):
                    yield e
        
        shape = (reduce(operator.mul, self.shape),)
        it = IterLayer(_iter(), shape=shape, dtype=self.dtype, length=self.length)
        if self.has_chunks:            
            return it.to_chunks(self.chunks_size, dtype=self.global_dtype)
        else:
{'fraction_non_alphanumeric': 0.0887582659808964, 'fraction_numerical': 0.003894195444526084, 'mean_word_length': 1.481042654028436, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 5, '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': '1628898', 'n_tokens_mistral': 3940, 'n_tokens_neox': 3874, 'n_words': 941}