#!/usr/bin/env python # -*- coding: utf-8 -*- # Copyright 1999-2022 Alibaba Group Holding Ltd. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import inspect from ...compat import six, reduce from ...models import TableSchema from ...utils import to_lower_str from ..utils import to_collection from .arithmetic import Equal from .core import Node, ExprDictionary from .expressions import CollectionExpr, TypedExpr, ProjectCollectionExpr, \ Column, BooleanSequenceExpr, SequenceExpr, Expr, Scalar, repr_obj, CallableColumn from .errors import ExpressionError class JoinCollectionExpr(CollectionExpr): __slots__ = '_how', '_left_suffix', '_right_suffix', '_column_origins', \ '_renamed_columns', '_column_conflict', '_mapjoin', '_skewjoin', \ '_skewjoin_values' _args = '_lhs', '_rhs', '_predicate' def _init(self, *args, **kwargs): self._init_attr('_left_suffix', None) self._init_attr('_right_suffix', None) self._init_attr('_column_origins', dict()) self._init_attr('_renamed_columns', dict()) self._init_attr('_column_conflict', False) super(JoinCollectionExpr, self)._init(*args, **kwargs) if not isinstance(self._lhs, CollectionExpr): raise TypeError( 'Can only join collection expressions, got %s for left expr.' % type(self._lhs) ) if not isinstance(self._rhs, CollectionExpr): raise TypeError( 'Can only join collection expressions, got %s for right expr.' % type(self._rhs) ) if self._rhs is self._lhs: self._rhs = self._rhs.view() if isinstance(self._lhs, JoinCollectionExpr) and \ (self._rhs is self._lhs._lhs or self._rhs is self._lhs._rhs): self._rhs = self._rhs.view() if self._left_suffix is None and self._right_suffix is None: overlaps = set(self._lhs.schema.names).intersection(self._rhs.schema.names) if len(overlaps) > 0: raise ValueError( 'Column conflict exists in join, overlap columns: %s' % ','.join(overlaps) ) self._set_schema() self._validate_predicates(self._predicate) self._how = self._how.upper() def _defunc(self, field): if inspect.isfunction(field): if six.get_function_code(field).co_argcount == 1: return field(self) else: return field(self._lhs, self._rhs) return field def _get_child(self, expr): while isinstance(expr, JoinProjectCollectionExpr): expr = expr.input return expr def _is_column(self, col, expr): return col.input is expr or col.input is self._get_child(expr) def _get_fields(self, fields, ret_raw_fields=False): selects = [] raw_selects = [] for field in fields: field = self._defunc(field) if isinstance(field, CollectionExpr): if any(c is self for c in field.children()) or \ any(c is self._get_child(self._lhs) for c in field.children()) or \ any(c is self._get_child(self._rhs) for c in field.children()): selects.extend(self._get_fields(field._project_fields)) elif field is self: selects.extend(self._get_fields(self._schema.names)) elif field is self._get_child(self._lhs): fields = [self._renamed_columns.get(n, [n])[0] for n in field.schema.names] selects.extend(self._get_fields(fields)) elif field is self._get_child(self._rhs): fields = [self._renamed_columns.get(n, [None, n])[1] for n in field.schema.names] selects.extend(self._get_fields(fields)) else: selects.extend(self._get_fields(field._fetch_fields())) raw_selects.append(field) else: select = self._get_field(field) selects.append(select) raw_selects.append(select) if ret_raw_fields: return selects, raw_selects return selects def _get_field(self, field): field = self._defunc(field) if isinstance(field, six.string_types): if field not in self._schema: raise ValueError('Field(%s) does not exist' % field) cls = Column if callable(getattr(type(self), field, None)): cls = CallableColumn return cls(self, _name=field, _data_type=self._schema[field].type) root = field has_path = False for expr in root.traverse(top_down=True, unique=True, stop_cond=lambda x: isinstance(x, Column) or x is self): if isinstance(expr, Column): if self._is_column(expr, self): has_path = True continue if self._is_column(expr, self._lhs): has_path = True idx = 0 elif self._is_column(expr, self._rhs): has_path = True idx = 1 elif expr.input._id == self._get_child(self._rhs)._id: # In case that the right collection has been copied expr.substitute(expr.input, self._get_child(self._rhs)) has_path = True idx = 1 elif isinstance(self._lhs, JoinCollectionExpr): try: expr = self._lhs._get_field(expr) except ExpressionError: continue has_path = True idx = 0 elif isinstance(self._rhs, JoinCollectionExpr): try: expr = self._rhs._get_field(expr) except ExpressionError: continue has_path = True idx = 1 else: continue name = expr.source_name if name in self._renamed_columns: name = self._renamed_columns[name][idx] to_sub = self._get_field(name) if expr.is_renamed(): to_sub = to_sub.rename(expr.name) to_sub.copy_to(expr) if isinstance(field, SequenceExpr) and not has_path: raise ExpressionError('field must come from Join collection ' 'or its left and right child collection: %s' % repr_obj(field)) return root def origin_collection(self, column_name): idx, name = self._column_origins[column_name] return [self._lhs, self._rhs][idx], name @property def node_name(self): return self.__class__.__name__ def iter_args(self): for it in zip(['collection(left)', 'collection(right)', 'on'], self.args): yield it @property def column_conflict(self): return self._column_conflict def accept(self, visitor): visitor.visit_join(self) def _get_non_suffixes_fields(self): return set() def _get_predicate_fields(self): predicate_fields = set() if not self._predicate: return predicate_fields for p in self._predicate: if isinstance(p, six.string_types): predicate_fields.add(p) elif isinstance(p, (tuple, Equal)): if isinstance(p, Equal): p = p.lhs, p.rhs if isinstance(p, tuple) and len(p) != 2: continue left_name = None if isinstance(p[0], six.string_types): left_name = p[0] elif isinstance(p[0], Column) and not p[0].is_renamed(): left_name = p[0].name if left_name is None: continue right_name = None if isinstance(p[1], six.string_types): right_name = p[1] elif isinstance(p[1], Column) and not p[1].is_renamed(): right_name = p[1].name if left_name == right_name: predicate_fields.add(left_name) return predicate_fields def _set_schema(self): names, typos = [], [] non_suffixes_fields = self._get_non_suffixes_fields() for col in self._lhs.schema.columns: name = col.name if to_lower_str(col.name) in self._rhs.schema._name_indexes: self._column_conflict = True if to_lower_str(col.name) in self._rhs.schema._name_indexes and \ col.name not in non_suffixes_fields: name = '%s%s' % (col.name, self._left_suffix) self._renamed_columns[col.name] = (name,) names.append(name) typos.append(col.type) self._column_origins[name] = 0, col.name for col in self._rhs.schema.columns: name = col.name if to_lower_str(col.name) in self._lhs.schema._name_indexes: self._column_conflict = True if to_lower_str(col.name) in self._lhs.schema._name_indexes and \ col.name not in non_suffixes_fields: name = '%s%s' % (col.name, self._right_suffix) self._renamed_columns[col.name] = \ self._renamed_columns[col.name][0], name if name in non_suffixes_fields: continue names.append(name) typos.append(col.type) self._column_origins[name] = 1, col.name if issubclass(type(self._lhs.schema), TableSchema): schema_type = type(self._lhs.schema) else: schema_type = type(self._rhs.schema) self._schema = schema_type.from_lists(names, typos) def _validate_equal(self, equal_expr): # FIXME: sometimes may be wrong, e.g. t3 = t1.join(t2, 'name')); t4.join(t3, t4.id == t3.id) return (equal_expr.lhs.is_ancestor(self._get_child(self._lhs)) and equal_expr.rhs.is_ancestor(self._get_child(self._rhs))) or \ (equal_expr.lhs.is_ancestor(self._get_child(self._rhs)) and equal_expr.rhs.is_ancestor(self._get_child(self._lhs))) def _reverse_equal(self, equal_expr): if equal_expr.lhs.is_ancestor(self._get_child(self._rhs)) and \ equal_expr.rhs.is_ancestor(self._get_child(self._lhs)): # the equal's left side is on the right collection and vise versa equal_expr._rhs, equal_expr._lhs = equal_expr._lhs, equal_expr._rhs def _validate_predicates(self, predicates): if predicates is None: return is_validate = False subs = [] if self._mapjoin: is_validate = True for p in predicates: if (isinstance(p, tuple) and len(p) == 2) or isinstance(p, six.string_types): if isinstance(p, six.string_types): left_name, right_name = p, p else: left_name, right_name = p left_col = self._lhs._get_field(left_name) right_col = self._rhs._get_field(right_name) if not left_col.is_ancestor(self._lhs) or not right_col.is_ancestor(self._rhs): raise ExpressionError('Invalid predicate: {0!s}'.format(repr_obj(p))) subs.append(left_col == right_col) is_validate = True elif isinstance(p, BooleanSequenceExpr): if not is_validate: it = (expr for expr in p.traverse(top_down=True, unique=True) if isinstance(expr, Equal)) while not is_validate: try: equal_expr = next(it) validate = self._validate_equal(equal_expr) if validate: is_validate = True break except StopIteration: break else: if not is_validate: raise ExpressionError('Invalid predicate: {0!s}'.format(repr_obj(p))) if not is_validate: raise ExpressionError('Invalid predicate: no validate predicate assigned') for p in predicates: if isinstance(p, BooleanSequenceExpr): subs.append(p) if len(subs) == 0: self._predicate = None else: self._predicate = subs def _merge_joined_fields(self, merge_columns): if not merge_columns: return self predicate_fields = self._get_predicate_fields() if not predicate_fields: raise ValueError('No fields in predicate. Cannot merge columns.') src_map = self._column_origins rename_map = dict() for name, src in six.iteritems(src_map): src_id, src_name = src if src_name not in predicate_fields: continue if src_name not in rename_map: rename_map[src_name] = [None, None] rename_map[src_name][src_id] = name if merge_columns in ('auto', 'left', 'right') or (isinstance(merge_columns, bool) and merge_columns): merge_columns = dict((k, merge_columns) for k in six.iterkeys(rename_map)) if isinstance(merge_columns, six.string_types): merge_columns = {merge_columns: 'auto'} if isinstance(merge_columns, list): merge_columns = dict((k, 'auto') for k in merge_columns) excludes = set() for col, action in six.iteritems(merge_columns): if col not in rename_map: raise ValueError('Column {0} not exists in join predicate.'.format(col)) if isinstance(action, bool) and action: merge_columns[col] = 'auto' else: merge_columns[col] = action.lower() excludes.update(rename_map[col]) selects = [] merged = set() for col in self.schema.names: if col not in excludes: selects.append(self[col]) else: src_name = src_map[col][1] if src_name in merged: continue merged.add(src_name) left_name, right_name = rename_map[src_name] left_col = self[left_name] right_col = self[right_name] if merge_columns[src_name] == 'auto': selects.append(left_col.isnull().ifelse(right_col, left_col).rename(src_name)) elif merge_columns[src_name] == 'left': selects.append(left_col.rename(src_name)) elif merge_columns[src_name] == 'right': selects.append(right_col.rename(src_name)) selected = self.select(*selects) return JoinFieldMergedCollectionExpr(_input=self, _fields=selected._fields, _schema=selected._schema, _rename_map=rename_map) class InnerJoin(JoinCollectionExpr): def _init(self, *args, **kwargs): self._how = 'INNER' super(InnerJoin, self)._init(*args, **kwargs) def _get_non_suffixes_fields(self): return self._get_predicate_fields() class LeftJoin(JoinCollectionExpr): def _init(self, *args, **kwargs): self._how = 'LEFT OUTER' super(LeftJoin, self)._init(*args, **kwargs) class RightJoin(JoinCollectionExpr): def _init(self, *args, **kwargs): self._how = 'RIGHT OUTER' super(RightJoin, self)._init(*args, **kwargs) class OuterJoin(JoinCollectionExpr): def _init(self, *args, **kwargs): self._how = 'FULL OUTER' super(OuterJoin, self)._init(*args, **kwargs) class JoinFieldMergedCollectionExpr(ProjectCollectionExpr): __slots__ = '_rename_map', def _get_fields(self, fields, ret_raw_fields=False): selects = [] raw_selects = [] joined_expr = self._input for field in fields: field = self._defunc(field) if isinstance(field, CollectionExpr): if any(c is self for c in field.children()) or \ any(c is joined_expr._lhs for c in field.children()) or \ any(c is joined_expr._rhs for c in field.children()): selects.extend(self._get_fields(field._project_fields)) elif field is self: selects.extend(self._get_fields(self._schema.names)) elif field is joined_expr._lhs: fields = [joined_expr._renamed_columns.get(n, [n])[0] if n not in self._rename_map else n for n in field.schema.names] selects.extend(self._get_fields(fields)) elif field is joined_expr._rhs: fields = [joined_expr._renamed_columns.get(n, [None, n])[1] if n not in self._rename_map else n for n in field.schema.names] selects.extend(self._get_fields(fields)) else: selects.extend(super(JoinFieldMergedCollectionExpr, self)._get_fields(field)) raw_selects.append(field) else: select = self._get_field(field) selects.append(select) raw_selects.append(select) if ret_raw_fields: return selects, raw_selects return selects def _get_field(self, field): field = self._defunc(field) if isinstance(field, six.string_types): if field not in self._schema: raise ValueError('Field(%s) does not exist' % field) cls = Column if callable(getattr(type(self), field, None)): cls = CallableColumn return cls(self, _name=field, _data_type=self._schema[field].type) joined_expr = self._input root = field has_path = False for expr in root.traverse(top_down=True, unique=True, stop_cond=lambda x: isinstance(x, Column) or x is self): if isinstance(expr, Column): if expr.input is self or expr.input is joined_expr: has_path = True continue if expr.input is joined_expr._lhs: has_path = True idx = 0 elif expr.input is joined_expr._rhs: has_path = True idx = 1 elif isinstance(joined_expr._lhs, JoinCollectionExpr): try: expr = joined_expr._lhs._get_field(expr) except ExpressionError: continue has_path = True idx = 0 elif isinstance(joined_expr._rhs, JoinCollectionExpr): try: expr = joined_expr._rhs._get_field(expr) except ExpressionError: continue has_path = True idx = 1 else: continue name = expr.source_name if name not in self._rename_map and name in joined_expr._renamed_columns: name = joined_expr._renamed_columns[name][idx] to_sub = self._get_field(name) if expr.is_renamed(): to_sub = to_sub.rename(expr.name) to_sub.copy_to(expr) if isinstance(field, SequenceExpr) and not has_path: raise ExpressionError('field must come from Join collection ' 'or its left and right child collection: %s' % repr_obj(field)) return root class JoinProjectCollectionExpr(ProjectCollectionExpr): """ Only for analyzer, project join should generate normal `ProjectCollectionExpr`. """ __slots__ = () _join_dict = { 'INNER': InnerJoin, 'LEFT': LeftJoin, 'RIGHT': RightJoin, 'OUTER': OuterJoin } def _make_different_sources(left, right, predicate=None): # TODO: move to analyzer, do it before analyze and optimize exprs = ExprDictionary() for n in left.traverse(unique=True): exprs[n] = True subs = ExprDictionary() if getattr(right, '_proxy', None) is not None: right = right._proxy dag = right.to_dag(copy=False, validate=False) for n in dag.traverse(): if n in exprs: copied = subs.get(n, n.copy()) for p in dag.successors(n): if p in exprs and p not in subs: subs[p] = p.copy() subs.get(p, p).substitute(n, copied) subs[n] = copied if predicate and n is right: for p in predicate: if not isinstance(p, Expr): continue p_dag = p.to_dag(copy=False, validate=False) for p_n in p_dag.traverse(top_down=True): if p_n is right: succs = list(s for s in p_dag.successors(p_n) if s not in exprs) p_dag.substitute(right, copied, parents=succs) return left, subs.get(right, right) def join(left, right, on=None, how='inner', suffixes=('_x', '_y'), mapjoin=False, skewjoin=False): """ Join two collections. If `on` is not specified, we will find the common fields of the left and right collection. `suffixes` means that if column names conflict, the suffixes will be added automatically. For example, both left and right has a field named `col`, there will be col_x, and col_y in the joined collection. :param left: left collection :param right: right collection :param on: fields to join on :param how: 'inner', 'left', 'right', or 'outer' :param suffixes: when name conflict, the suffix will be added to both columns. :param mapjoin: set use mapjoin or not, default value False. :param skewjoin: set use of skewjoin or not, default value False. Can specify True if the collection is skew, or a list specifying columns with skew values, or a list of dicts specifying skew combinations. :return: collection :Example: >>> df.dtypes.names ['name', 'id'] >>> df2.dtypes.names ['name', 'id1'] >>> df.join(df2) >>> df.join(df2, on='name') >>> df.join(df2, on=('id', 'id1')) >>> df.join(df2, on=['name', ('id', 'id1')]) >>> df.join(df2, on=[df.name == df2.name, df.id == df2.id1]) >>> df.join(df2, mapjoin=False) >>> df.join(df2, skewjoin=True) >>> df.join(df2, skewjoin=["c0", "c1"]) >>> df.join(df2, skewjoin=[{"c0": 1, "c1": "2"}, {"c0": 3, "c1": "4"}]) """ if mapjoin and skewjoin: raise TypeError("Cannot specify mapjoin and skewjoin at the same time") if isinstance(left, TypedExpr): left = to_collection(left) if isinstance(right, TypedExpr): right = to_collection(right) if on is None: if not mapjoin: on = [name for name in left.schema.names if to_lower_str(name) in right.schema._name_indexes] if not on and len(left.schema) == 1 and len(right.schema) == 1: on = [(left.schema.names[0], right.schema.names[0])] skewjoin_values = None if isinstance(skewjoin, (dict, six.string_types)): skewjoin = [skewjoin] if isinstance(skewjoin, list): if ( all(isinstance(c, six.string_types) for c in skewjoin) and any(c not in right.schema.names for c in skewjoin) ): raise ValueError( "All columns specified in `skewjoin` need to exist in the right collection" ) elif ( all(isinstance(c, dict) for c in skewjoin) ): cols = sorted(skewjoin[0].keys()) cols_set = set(cols) if any(c not in right.schema.names for c in cols): raise ValueError( "All columns specified in `skewjoin` need to exist in the right collection" ) if any(cols_set != set(c.keys()) for c in skewjoin): raise ValueError("All values in `skewjoin` need to have same columns") skewjoin_values = [[d[c] for c in cols] for d in skewjoin] skewjoin = cols elif skewjoin and not isinstance(skewjoin, bool): raise TypeError("Cannot accept skewjoin type %s" % type(skewjoin)) if isinstance(suffixes, (tuple, list)) and len(suffixes) == 2: left_suffix, right_suffix = suffixes else: raise ValueError('suffixes must be a tuple or list with two elements, got %s' % suffixes) if not isinstance(on, list): on = [on, ] for i in range(len(on)): it = on[i] if inspect.isfunction(it): on[i] = it(left, right) left, right = _make_different_sources(left, right, on) try: return _join_dict[how.upper()]( _lhs=left, _rhs=right, _predicate=on, _left_suffix=left_suffix, _right_suffix=right_suffix, _mapjoin=mapjoin, _skewjoin=skewjoin, _skewjoin_values=skewjoin_values ) except KeyError: return JoinCollectionExpr( _lhs=left, _rhs=right, _predicate=on, _how=how, _left_suffix=left_suffix, _right_suffix=right_suffix, _mapjoin=mapjoin, _skewjoin=skewjoin, _skewjoin_values=skewjoin_values ) def inner_join(left, right, on=None, suffixes=('_x', '_y'), mapjoin=False, skewjoin=False): """ Inner join two collections. If `on` is not specified, we will find the common fields of the left and right collection. `suffixes` means that if column names conflict, the suffixes will be added automatically. For example, both left and right has a field named `col`, there will be col_x, and col_y in the joined collection. :param left: left collection :param right: right collection :param on: fields to join on :param suffixes: when name conflict, the suffixes will be added to both columns. :return: collection :Example: >>> df.dtypes.names ['name', 'id'] >>> df2.dtypes.names ['name', 'id1'] >>> df.inner_join(df2) >>> df.inner_join(df2, on='name') >>> df.inner_join(df2, on=('id', 'id1')) >>> df.inner_join(df2, on=['name', ('id', 'id1')]) >>> df.inner_join(df2, on=[df.name == df2.name, df.id == df2.id1]) """ return join(left, right, on, suffixes=suffixes, mapjoin=mapjoin, skewjoin=skewjoin) def left_join( left, right, on=None, suffixes=('_x', '_y'), mapjoin=False, merge_columns=None, skewjoin=False ): """ Left join two collections. If `on` is not specified, we will find the common fields of the left and right collection. `suffixes` means that if column names conflict, the suffixes will be added automatically. For example, both left and right has a field named `col`, there will be col_x, and col_y in the joined collection. :param left: left collection :param right: right collection :param on: fields to join on :param suffixes: when name conflict, the suffixes will be added to both columns. :param mapjoin: set use mapjoin or not, default value False. :param merge_columns: whether to merge columns with the same name into one column without suffix. If the value is True, columns in the predicate with same names will be merged, with non-null value. If the value is 'left' or 'right', the values of predicates on the left / right collection will be taken. You can also pass a dictionary to describe the behavior of each column, such as { 'a': 'auto', 'b': 'left' }. :return: collection :Example: >>> df.dtypes.names ['name', 'id'] >>> df2.dtypes.names ['name', 'id1'] >>> df.left_join(df2) >>> df.left_join(df2, on='name') >>> df.left_join(df2, on=('id', 'id1')) >>> df.left_join(df2, on=['name', ('id', 'id1')]) >>> df.left_join(df2, on=[df.name == df2.name, df.id == df2.id1]) """ joined = join(left, right, on, how='left', suffixes=suffixes, mapjoin=mapjoin, skewjoin=skewjoin) return joined._merge_joined_fields(merge_columns) def right_join( left, right, on=None, suffixes=('_x', '_y'), mapjoin=False, merge_columns=None, skewjoin=False ): """ Right join two collections. If `on` is not specified, we will find the common fields of the left and right collection. `suffixes` means that if column names conflict, the suffixes will be added automatically. For example, both left and right has a field named `col`, there will be col_x, and col_y in the joined collection. :param left: left collection :param right: right collection :param on: fields to join on :param suffixes: when name conflict, the suffixes will be added to both columns. :param mapjoin: set use mapjoin or not, default value False. :param merge_columns: whether to merge columns with the same name into one column without suffix. If the value is True, columns in the predicate with same names will be merged, with non-null value. If the value is 'left' or 'right', the values of predicates on the left / right collection will be taken. You can also pass a dictionary to describe the behavior of each column, such as { 'a': 'auto', 'b': 'left' }. :return: collection :Example: >>> df.dtypes.names ['name', 'id'] >>> df2.dtypes.names ['name', 'id1'] >>> df.right_join(df2) >>> df.right_join(df2, on='name') >>> df.right_join(df2, on=('id', 'id1')) >>> df.right_join(df2, on=['name', ('id', 'id1')]) >>> df.right_join(df2, on=[df.name == df2.name, df.id == df2.id1]) """ joined = join(left, right, on, how='right', suffixes=suffixes, mapjoin=mapjoin, skewjoin=skewjoin) return joined._merge_joined_fields(merge_columns) def outer_join( left, right, on=None, suffixes=('_x', '_y'), mapjoin=False, merge_columns=None, skewjoin=False ): """ Outer join two collections. If `on` is not specified, we will find the common fields of the left and right collection. `suffixes` means that if column names conflict, the suffixes will be added automatically. For example, both left and right has a field named `col`, there will be col_x, and col_y in the joined collection. :param left: left collection :param right: right collection :param on: fields to join on :param suffixes: when name conflict, the suffixes will be added to both columns. :param mapjoin: set use mapjoin or not, default value False. :param merge_columns: whether to merge columns with the same name into one column without suffix. If the value is True, columns in the predicate with same names will be merged, with non-null value. If the value is 'left' or 'right', the values of predicates on the left / right collection will be taken. You can also pass a dictionary to describe the behavior of each column, such as { 'a': 'auto', 'b': 'left' }. :return: collection :Example: >>> df.dtypes.names ['name', 'id'] >>> df2.dtypes.names ['name', 'id1'] >>> df.outer_join(df2) >>> df.outer_join(df2, on='name') >>> df.outer_join(df2, on=('id', 'id1')) >>> df.outer_join(df2, on=['name', ('id', 'id1')]) >>> df.outer_join(df2, on=[df.name == df2.name, df.id == df2.id1]) """ joined = join(left, right, on, how='outer', suffixes=suffixes, mapjoin=mapjoin, skewjoin=skewjoin) return joined._merge_joined_fields(merge_columns) CollectionExpr.join = join CollectionExpr.inner_join = inner_join CollectionExpr.left_join = left_join CollectionExpr.right_join = right_join CollectionExpr.outer_join = outer_join TypedExpr.join = join TypedExpr.inner_join = inner_join TypedExpr.left_join = left_join TypedExpr.right_join = right_join TypedExpr.outer_join = outer_join def _get_sequence_source_collection(expr): return next(it for it in expr.traverse(top_down=True, unique=True) if isinstance(it, CollectionExpr)) class UnionCollectionExpr(CollectionExpr): __slots__ = '_distinct', _args = '_lhs', '_rhs', node_name = 'Union' def _init(self, *args, **kwargs): super(UnionCollectionExpr, self)._init(*args, **kwargs) self._validate() self._schema = self._clean_schema() def _validate_collection_child(self): if self._lhs.schema.names == self._rhs.schema.names and self._lhs.schema.types == self._rhs.schema.types: return True elif set(self._lhs.schema.names) == set(self._rhs.schema.names) and set(self._lhs.schema.types) == set( self._rhs.schema.types): self._rhs = self._rhs[self._lhs.schema.names] return self._lhs.schema.types == self._rhs.schema.types else: return False def _validate(self): if isinstance(self._lhs, SequenceExpr): source_collection = _get_sequence_source_collection(self._lhs) self._lhs = source_collection[[self._lhs]] if isinstance(self._rhs, SequenceExpr): source_collection = _get_sequence_source_collection(self._rhs) self._rhs = source_collection[[self._rhs]] if isinstance(self._lhs, CollectionExpr) and isinstance(self._rhs, CollectionExpr): if not self._validate_collection_child(): raise ExpressionError('Table schemas must be equal to form union') else: raise ExpressionError('Both inputs should be collections or sequences.') def _clean_schema(self): return TableSchema.from_lists(self._lhs.schema.names, self._lhs.schema.types) def accept(self, visitor): return visitor.visit_union(self) def iter_args(self): for it in zip(['collection(left)', 'collection(right)'], self.args): yield it class ConcatCollectionExpr(CollectionExpr): _args = '_lhs', '_rhs', node_name = 'Concat' def _init(self, *args, **kwargs): super(ConcatCollectionExpr, self)._init(*args, **kwargs) self._schema = self._clean_schema() def iter_args(self): for it in zip(['collection(left)', 'collection(right)'], self.args): yield it def _clean_schema(self): return TableSchema.from_lists( self._lhs.schema.names + self._rhs.schema.names, self._lhs.schema.types + self._rhs.schema.types, ) @staticmethod def _get_sequence_source_collection(expr): return next(it for it in expr.traverse(top_down=True, unique=True) if isinstance(it, CollectionExpr)) @classmethod def validate_input(cls, *inputs): new_inputs = [] for i in inputs: if isinstance(i, SequenceExpr): source_collection = cls._get_sequence_source_collection(i) new_inputs.append(source_collection[[i]]) else: new_inputs.append(i) if any(not isinstance(i, CollectionExpr) for i in new_inputs): raise ExpressionError('Inputs should be collections or sequences.') unioned = reduce(lambda a, b: a | b, (set(i.schema.names) for i in inputs)) total_fields = sum((len(i.schema.names) for i in inputs)) if total_fields != len(unioned): raise ExpressionError('Column names in inputs should not collides with each other.') def accept(self, visitor): return visitor.visit_concat(self) def union(left, right, distinct=False): """ Union two collections. :param left: left collection :param right: right collection :param distinct: :return: collection :Example: >>> df['name', 'id'].union(df2['id', 'name']) """ left, right = _make_different_sources(left, right) return UnionCollectionExpr(_lhs=left, _rhs=right, _distinct=distinct) def __horz_concat(left, rights): from ..utils import to_collection left = to_collection(left) for right in rights: right = to_collection(right) left, right = _make_different_sources(left, right) left = ConcatCollectionExpr(_lhs=left, _rhs=right) return left def concat(left, rights, distinct=False, axis=0): """ Concat collections. :param left: left collection :param rights: right collections, can be a DataFrame object or a list of DataFrames :param distinct: whether to remove duplicate entries. only available when axis == 0 :param axis: when axis == 0, the DataFrames are merged vertically, otherwise horizontally. :return: collection Note that axis==1 can only be used under Pandas DataFrames or XFlow. :Example: >>> df['name', 'id'].concat(df2['score'], axis=1) """ from ..utils import to_collection if isinstance(rights, Node): rights = [rights, ] if not rights: raise ValueError('At least one DataFrame should be provided.') if axis == 0: for right in rights: left = union(left, right, distinct=distinct) return left else: rights = [to_collection(r) for r in rights] ConcatCollectionExpr.validate_input(left, *rights) if hasattr(left, '_xflow_concat'): return left._xflow_concat(rights) else: return __horz_concat(left, rights) def _drop(expr, data, axis=0, columns=None): """ Drop data from a DataFrame. :param expr: collection to drop data from :param data: data to be removed :param axis: 0 for deleting rows, 1 for columns. :param columns: columns of data to select, only useful when axis == 0 :return: collection :Example: >>> import pandas as pd >>> df1 = DataFrame(pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]})) >>> df2 = DataFrame(pd.DataFrame({'a': [2, 3], 'b': [5, 7]})) >>> df1.drop(df2) a b c 0 1 4 7 1 3 6 9 >>> df1.drop(df2, columns='a') a b c 0 1 4 7 >>> df1.drop(['a'], axis=1) b c 0 4 7 1 5 8 2 6 9 >>> df1.drop(df2, axis=1) c 0 7 1 8 2 9 """ from ..utils import to_collection expr = to_collection(expr) if axis == 0: if not isinstance(data, (CollectionExpr, SequenceExpr)): raise ExpressionError('data should be a collection or sequence when axis == 1.') data = to_collection(data) if columns is None: columns = [n for n in data.schema.names] if isinstance(columns, six.string_types): columns = [columns, ] data = data.select(*columns).distinct() drop_predicates = [data[n].isnull() for n in data.schema.names] return expr.left_join(data, on=columns, suffixes=('', '_dp')).filter(*drop_predicates) \ .select(*expr.schema.names) else: if isinstance(data, (CollectionExpr, SequenceExpr)): data = to_collection(data).schema.names return expr.exclude(data) def setdiff(left, *rights, **kwargs): """ Exclude data from a collection, like `except` clause in SQL. All collections involved should have same schema. :param left: collection to drop data from :param rights: collection or list of collections :param distinct: whether to preserve duplicate entries :return: collection :Examples: >>> import pandas as pd >>> df1 = DataFrame(pd.DataFrame({'a': [1, 2, 3, 3, 3], 'b': [1, 2, 3, 3, 3]})) >>> df2 = DataFrame(pd.DataFrame({'a': [1, 3], 'b': [1, 3]})) >>> df1.setdiff(df2) a b 0 2 2 1 3 3 2 3 3 >>> df1.setdiff(df2, distinct=True) a b 0 2 2 """ import time from ..utils import output distinct = kwargs.get('distinct', False) if isinstance(rights[0], list): rights = rights[0] cols = [n for n in left.schema.names] types = [n for n in left.schema.types] counter_col_name = 'exc_counter_%d' % int(time.time()) left = left[left, Scalar(1).rename(counter_col_name)] rights = [r[r, Scalar(-1).rename(counter_col_name)] for r in rights] unioned = left for r in rights: unioned = unioned.union(r) if distinct: aggregated = unioned.groupby(*cols).agg(**{counter_col_name: unioned[counter_col_name].min()}) return aggregated.filter(aggregated[counter_col_name] == 1).select(*cols) else: aggregated = unioned.groupby(*cols).agg(**{counter_col_name: unioned[counter_col_name].sum()}) @output(cols, types) def exploder(row): import sys irange = xrange if sys.version_info[0] < 3 else range for _ in irange(getattr(row, counter_col_name)): yield row[:-1] return aggregated.map_reduce(mapper=exploder).select(*cols) def intersect(left, *rights, **kwargs): """ Calc intersection among datasets, :param left: collection :param rights: collection or list of collections :param distinct: whether to preserve duolicate entries :return: collection :Examples: >>> import pandas as pd >>> df1 = DataFrame(pd.DataFrame({'a': [1, 2, 3, 3, 3], 'b': [1, 2, 3, 3, 3]})) >>> df2 = DataFrame(pd.DataFrame({'a': [1, 3, 3], 'b': [1, 3, 3]})) >>> df1.intersect(df2) a b 0 1 1 1 3 3 2 3 3 >>> df1.intersect(df2, distinct=True) a b 0 1 1 1 3 3 """ import time from ..utils import output distinct = kwargs.get('distinct', False) if isinstance(rights[0], list): rights = rights[0] cols = [n for n in left.schema.names] types = [n for n in left.schema.types] collections = (left, ) + rights idx_col_name = 'idx_%d' % int(time.time()) counter_col_name = 'exc_counter_%d' % int(time.time()) collections = [c[c, Scalar(idx).rename(idx_col_name)] for idx, c in enumerate(collections)] unioned = reduce(lambda a, b: a.union(b), collections) src_agg = unioned.groupby(*(cols + [idx_col_name])) \ .agg(**{counter_col_name: unioned.count()}) aggregators = { idx_col_name: src_agg[idx_col_name].nunique(), counter_col_name: src_agg[counter_col_name].min(), } final_agg = src_agg.groupby(*cols).agg(**aggregators) final_agg = final_agg.filter(final_agg[idx_col_name] == len(collections)) if distinct: return final_agg.filter(final_agg[counter_col_name] > 0).select(*cols) else: @output(cols, types) def exploder(row): import sys irange = xrange if sys.version_info[0] < 3 else range for _ in irange(getattr(row, counter_col_name)): yield row[:-2] return final_agg.map_reduce(mapper=exploder).select(*cols) CollectionExpr.union = union SequenceExpr.union = union CollectionExpr.__horz_concat = __horz_concat SequenceExpr.__horz_concat = __horz_concat CollectionExpr.concat = concat SequenceExpr.concat = concat CollectionExpr.drop = _drop SequenceExpr.drop = _drop CollectionExpr.setdiff = setdiff SequenceExpr.setdiff = setdiff CollectionExpr.except_ = setdiff SequenceExpr.except_ = setdiff CollectionExpr.intersect = intersect SequenceExpr.intersect = intersect