Column Assignment
| 1959 | |
| 1960 | |
| 1961 | class Assign(Elemwise): |
| 1962 | """Column Assignment""" |
| 1963 | |
| 1964 | _parameters = ["frame"] |
| 1965 | operation = staticmethod(assign) |
| 1966 | |
| 1967 | @functools.cached_property |
| 1968 | def unique_partition_mapping_columns_from_shuffle(self): |
| 1969 | keys = set(self.keys) |
| 1970 | return { |
| 1971 | col |
| 1972 | for col in self.frame.unique_partition_mapping_columns_from_shuffle |
| 1973 | if not isinstance(col, tuple) |
| 1974 | and col not in keys |
| 1975 | or not set(col).intersection(keys) |
| 1976 | } |
| 1977 | |
| 1978 | @functools.cached_property |
| 1979 | def keys(self): |
| 1980 | return self.operands[1::2] |
| 1981 | |
| 1982 | @functools.cached_property |
| 1983 | def vals(self): |
| 1984 | return self.operands[2::2] |
| 1985 | |
| 1986 | @functools.cached_property |
| 1987 | def _meta(self): |
| 1988 | args = [op._meta if isinstance(op, Expr) else op for op in self._args] |
| 1989 | return make_meta(self.operation(*args, **self._kwargs)) |
| 1990 | |
| 1991 | def _tree_repr_argument_construction(self, i, op, header): |
| 1992 | if i == 0: |
| 1993 | return super()._tree_repr_argument_construction(i, op, header) |
| 1994 | if i % 2 == 1: |
| 1995 | sep = "" if i == 1 else "," |
| 1996 | header += f"{sep} {repr(op)[1:-1]}=" |
| 1997 | else: |
| 1998 | header += f"{op!r}" |
| 1999 | return header |
| 2000 | |
| 2001 | def _node_label_args(self): |
| 2002 | return self.operands |
| 2003 | |
| 2004 | def _remove_common_columns(self, other): |
| 2005 | if set(self.keys) & set(other.keys): |
| 2006 | keys = set(self.keys) |
| 2007 | operands = [[k, v] for k, v in zip(other.keys, other.vals) if k not in keys] |
| 2008 | return [other.frame] + list(flatten(operands)) + self.operands[1:] |
| 2009 | else: |
| 2010 | return other.operands + self.operands[1:] |
| 2011 | |
| 2012 | def _simplify_down(self): |
| 2013 | if isinstance(self.frame, Assign): |
| 2014 | if self._check_for_previously_created_column(self.frame): |
| 2015 | # don't squash if we are using a column that was previously created |
| 2016 | return |
| 2017 | return Assign(*self._remove_common_columns(self.frame)) |
| 2018 | elif isinstance(self.frame, Projection) and isinstance( |