def wrapper(self, *args): # It is possible for the function `f` takes other arguments than Spark Column. # To cover this case, explicitly check if the argument is Koalas Series and # extract Spark Column. For other arguments, they are used as are. cols = [arg for arg in args if isinstance(arg, IndexOpsMixin)] if all(same_anchor(self, col) for col in cols): # Same DataFrame anchors args = [arg.spark.column if isinstance(arg, IndexOpsMixin) else arg for arg in args] scol = f(self.spark.column, *args) scol = booleanize_null(self.spark.column, scol, f) kser = self._with_new_scol(scol) else: # Different DataFrame anchors def apply_func(this_column, *that_columns): scol = f(this_column, *that_columns) return booleanize_null(this_column, scol, f) kser = align_diff_series(apply_func, self, *args, how="full") if not all(self.name == col.name for col in cols): kser = kser.rename() return kser
def wrapper(self, *args): # It is possible for the function `f` takes other arguments than Spark Column. # To cover this case, explicitly check if the argument is Koalas Series and # extract Spark Column. For other arguments, they are used as are. cols = [arg for arg in args if isinstance(arg, IndexOpsMixin)] if all(self._kdf is col._kdf for col in cols): # Same DataFrame anchors args = [arg._scol if isinstance(arg, IndexOpsMixin) else arg for arg in args] scol = f(self._scol, *args) return self._with_new_scol(scol) else: # Different DataFrame anchors def apply_func(this_column, *that_columns): return f(this_column, *that_columns) return align_diff_series(apply_func, self, *args, how="full")
def wrapper(self, *args): # It is possible for the function `f` takes other arguments than Spark Column. # To cover this case, explicitly check if the argument is Koalas Series and # extract Spark Column. For other arguments, they are used as are. cols = [arg for arg in args if isinstance(arg, IndexOpsMixin)] if all(self._kdf is col._kdf for col in cols): # Same DataFrame anchors args = [ arg._scol if isinstance(arg, IndexOpsMixin) else arg for arg in args ] scol = f(self._scol, *args) # check if `f` is a comparison operator comp_ops = ['eq', 'ne', 'lt', 'le', 'ge', 'gt'] is_comp_op = any( f == getattr(spark.Column, '__{}__'.format(comp_op)) for comp_op in comp_ops) if is_comp_op: filler = f == spark.Column.__ne__ scol = F.when(scol.isNull(), filler).otherwise(scol) elif f == spark.Column.__or__: scol = F.when(self._scol.isNull() | scol.isNull(), False).otherwise(scol) elif f == spark.Column.__and__: scol = F.when(scol.isNull(), False).otherwise(scol) return self._with_new_scol(scol) else: # Different DataFrame anchors def apply_func(this_column, *that_columns): return f(this_column, *that_columns) return align_diff_series(apply_func, self, *args, how="full")