def rename(columns_old_new=None, func=None): """" Changes the name of a column(s) dataFrame. :param columns_old_new: List of tuples. Each tuple has de following form: (oldColumnName, newColumnName). :param func: can be lower, upper or any string transformation function """ df = self # Apply a transformation function if is_function(func): exprs = [F.col(c).alias(func(c)) for c in df.columns] df = df.select(exprs) elif is_list_of_tuples(columns_old_new): # Check that the 1st element in the tuple is a valid set of columns validate_columns_names(self, columns_old_new) for c in columns_old_new: old_col_name = c[0] if is_str(old_col_name): df = df.withColumnRenamed(old_col_name, c[1]) elif is_int(old_col_name): df = df.withColumnRenamed(self.schema.names[old_col_name], c[1]) return df
def replace(columns, search_and_replace=None, value=None, regex=None): """ Replace a value or a list of values by a specified string :param columns: '*', list of columns names or a single column name. :param search_and_replace: values to look at to be replaced :param value: new value to replace the old one :param regex: :return: """ replace = None search = None if is_list_of_tuples(search_and_replace): params = list(zip(*search_and_replace)) search = list(params[0]) replace = list(params[1]) elif is_list(search_and_replace): search = search_and_replace replace = value elif is_one_element(search_and_replace): search = val_to_list(search_and_replace) replace = value if regex: search = search_and_replace replace = value # if regex or normal replace we use regexp or replace functions # TODO check if .contains can be used instead of regexp def func_regex(_df, _col_name, _search, _replace): return _df.withColumn( c, F.regexp_replace(_col_name, _search, _replace)) def func_replace(_df, _col_name, _search, _replace): data_type = self.cols.dtypes(_col_name) _search = [PYTHON_TYPES_[data_type](s) for s in _search] _df = _df.replace(_search, _replace, _col_name) return _df if regex: func = func_regex else: func = func_replace df = self columns = parse_columns(self, columns, filter_by_column_dtypes="string") for c in columns: df = func(df, c, search, replace) return df
def data_frame(cols=None, rows=None, infer_schema=True, pdf=None): """ Helper to create a Spark dataframe: :param cols: List of Tuple with name, data type and a flag to accept null :param rows: List of Tuples with the same number and types that cols :param infer_schema: Try to infer the schema data type. :param pdf: a pandas dataframe :return: Dataframe """ if is_(pdf, pd.DataFrame): result = Spark.instance.spark.createDataFrame(pdf) else: specs = [] # Process the rows if not is_list_of_tuples(rows): rows = [(i, ) for i in rows] # Process the columns for c, r in zip(cols, rows[0]): # Get columns name if is_one_element(c): col_name = c if infer_schema is True: var_type = infer(r) else: var_type = StringType() nullable = True elif is_tuple(c): # Get columns data type col_name = c[0] var_type = get_spark_dtypes_object(c[1]) count = len(c) if count == 2: nullable = True elif count == 3: nullable = c[2] # If tuple has not the third param with put it to true to accepts Null in columns specs.append([col_name, var_type, nullable]) struct_fields = list(map(lambda x: StructField(*x), specs)) result = Spark.instance.spark.createDataFrame( rows, StructType(struct_fields)) return result
def append(row): """ Append a row at the end of a dataframe :param row: List of values or tuples to be appended :return: Spark DataFrame """ df = self columns = [str(i) for i in range(df.cols.count())] if not is_list_of_tuples(row): row = [tuple(row)] new_row = op.Create.df(columns, row) return df.union(new_row)
def validate_columns_names(df, col_names, index=0): """ Check if a string or list of string are valid dataframe columns :param df: :param col_names: :param index: :return: """ columns = val_to_list(col_names) if is_list_of_tuples(columns): columns = [c[index] for c in columns] # Remove duplicates in the list if is_list_of_strings(columns): columns = set(columns) check_for_missing_columns(df, columns) return True
def append(cols_values=None): """ Append a column or a Dataframe to a Dataframe :param cols_values: New Column Names and data values :type cols_values: List of tuples :return: """ df_result = None # Append a dataframe if is_list_of_dataframes(cols_values): dfs = cols_values dfs.insert(0, self) df_result = append(dfs, like="columns") elif is_list_of_tuples(cols_values): df_result = self for c in cols_values: col_name = c[0] value = c[1] df_result = df_result.cols.append(col_name, value) return df_result
def create(self, df, func, suffix=None, output="df", *args, **kwargs): """ This is a helper function that output python tests for Spark Dataframes. :param df: Spark Dataframe :param suffix: The create method will try to create a test function with the func param given. If you want to test a function with different params you can use suffix. :param func: Spark dataframe function to be tested :param output: can be a 'df' or a 'json' :param args: Arguments to be used in the function :param kwargs: Keyword arguments to be used in the functions :return: """ buffer = [] def add_buffer(value): buffer.append("\t" + value) if suffix is None: suffix = "" else: suffix = "_" + suffix # Create func test name. If is None we just test the create.df function a not transform the data frame in # any way if func is None: func_test_name = "test_" + "create_df" + suffix + "()" else: func_test_name = "test_" + func.replace(".", "_") + suffix + "()" print("Creating {test} test function...".format(test=func_test_name)) logging.info(func_test_name) add_buffer("@staticmethod\n") add_buffer("def " + func_test_name + ":\n") if df is not None: source_df = "\tsource_df=op.create.df(" + df.export() + ")\n" df_func = df add_buffer(source_df) else: df_func = self.df # Process simple arguments _args = [] for v in args: if is_str(v): _args.append("'" + v + "'") elif is_numeric(v): _args.append(str(v)) elif is_list(v): if is_list_of_strings(v): lst = ["'" + x + "'" for x in v] elif is_list_of_numeric(v): lst = [str(x) for x in v] elif is_list_of_tuples(v): lst = [str(x) for x in v] _args.append('[' + ','.join(lst) + ']') _args = ','.join(_args) _kwargs = [] # print(_args) # Process keywords arguments for k, v in kwargs.items(): if is_str(v): v = "'" + v + "'" _kwargs.append(k + "=" + str(v)) # Separator if we have positional and keyword arguments separator = "" if (not is_list_empty(args)) & (not is_list_empty(kwargs)): separator = "," if func is None: add_buffer("\tactual_df = source_df\n") else: add_buffer("\tactual_df = source_df." + func + "(" + _args + separator + ','.join(_kwargs) + ")\n") # Apply function to the dataframe if func is None: df_result = self.op.create.df(*args, **kwargs) else: for f in func.split("."): df_func = getattr(df_func, f) df_result = df_func(*args, **kwargs) if output == "df": expected = "\texpected_df = op.create.df(" + df_result.export( ) + ")\n" elif output == "json": if is_str(df_result): df_result = "'" + df_result + "'" else: df_result = str(df_result) expected = "\texpected_value =" + df_result + "\n" add_buffer(expected) if output == "df": add_buffer( "\tassert (expected_df.collect() == actual_df.collect())\n") elif output == "json": add_buffer("\tassert (expected_value == actual_df)\n") return "".join(buffer)
def create(self, df, func, suffix=None, output="df", *args, **kwargs): """ This is a helper function that output python tests for Spark Dataframes. :param df: Spark Dataframe :param suffix: The create method will try to create a test function with the func param given. If you want to test a function with different params you can use suffix. :param func: Spark dataframe function to be tested :param output: can be a 'df' or a 'json' :param args: Arguments to be used in the function :param kwargs: Keyword arguments to be used in the functions :return: """ buffer = [] def add_buffer(value): buffer.append("\t" + value) if suffix is None: suffix = "" else: suffix = "_" + suffix # Create func test name. If is None we just test the create.df function a not transform the data frame in # any way if func is None: func_test_name = "test_" + "create_df" + suffix + "()" else: func_test_name = "test_" + func.replace(".", "_") + suffix + "()" print("Creating {test} test function...".format(test=func_test_name)) logger.print(func_test_name) add_buffer("@staticmethod\n") add_buffer("def " + func_test_name + ":\n") source = "source_df" if df is None: # Use the main df df_func = self.df elif isinstance(df, pyspark.sql.dataframe.DataFrame): source_df = "\tsource_df=op.create.df(" + df.export() + ")\n" df_func = df add_buffer(source_df) else: # TODO: op is not supposed to be hardcoded source = "op" df_func = df # Process simple arguments _args = [] for v in args: if is_str(v): _args.append("'" + v + "'") elif is_numeric(v): _args.append(str(v)) elif is_list(v): if is_list_of_strings(v): lst = ["'" + x + "'" for x in v] elif is_list_of_numeric(v): lst = [str(x) for x in v] elif is_list_of_tuples(v): lst = [str(x) for x in v] _args.append('[' + ','.join(lst) + ']') elif is_function(v): _args.append(v.__qualname__) # else: # import marshal # code_string = marshal.dumps(v.__code__) # add_buffer("\tfunction = '" + code_string + "'\n") # import marshal, types # # code = marshal.loads(code_string) # func = types.FunctionType(code, globals(), "some_func_name") _args = ','.join(_args) _kwargs = [] # print(_args) # Process keywords arguments for k, v in kwargs.items(): if is_str(v): v = "'" + v + "'" _kwargs.append(k + "=" + str(v)) # Separator if we have positional and keyword arguments separator = "" if (not is_list_empty(args)) & (not is_list_empty(kwargs)): separator = "," if func is None: add_buffer("\tactual_df = source_df\n") else: add_buffer("\tactual_df =" + source + "." + func + "(" + _args + separator + ','.join(_kwargs) + ")\n") # Apply function to the dataframe if func is None: df_result = self.op.create.df(*args, **kwargs) else: # Here we construct the method to be applied to the source object for f in func.split("."): df_func = getattr(df_func, f) df_result = df_func(*args, **kwargs) if output == "df": expected = "\texpected_df = op.create.df(" + df_result.export( ) + ")\n" elif output == "json": if is_str(df_result): df_result = "'" + df_result + "'" else: df_result = str(df_result) add_buffer("\tactual_df =json_enconding(actual_df)\n") expected = "\texpected_value =json_enconding(" + df_result + ")\n" else: expected = "\t\n" add_buffer(expected) if output == "df": add_buffer( "\tassert (expected_df.collect() == actual_df.collect())\n") elif output == "json": add_buffer("\tassert (expected_value == actual_df)\n") return "".join(buffer)
def parse_columns(df, cols_args, get_args=False, is_regex=None, filter_by_column_dtypes=None, accepts_missing_cols=False): """ Return a list of columns and check that columns exists in the dataframe Accept '*' as parameter in which case return a list of all columns in the dataframe. Also accept a regex. If a list of tuples return to list. The first element is the columns name the others element are params. This params can me used to create custom transformation functions. You can find and example in cols().cast() :param df: Dataframe in which the columns are going to be checked :param cols_args: Accepts * as param to return all the string columns in the dataframe :param get_args: :param is_regex: Use True is col_attrs is a regex :param filter_by_column_dtypes: :param accepts_missing_cols: if true not check if column exist in the dataframe :return: A list of columns string names """ cols = None attrs = None # ensure that cols_args is a list # cols_args = val_to_list(cols_args) # if columns value is * get all dataframes columns if is_regex is True: r = re.compile(cols_args[0]) cols = list(filter(r.match, df.columns)) elif cols_args == "*" or cols_args is None: cols = df.columns # In case we have a list of tuples we use the first element of the tuple is taken as the column name # and the rest as params. We can use the param in a custom function as follow # def func(attrs): attrs return (1,2) and (3,4) # return attrs[0] + 1 # df.cols().apply([('col_1',1,2),('cols_2', 3 ,4)], func) # Verify if we have a list with tuples elif is_tuple(cols_args) or is_list_of_tuples(cols_args): cols_args = val_to_list(cols_args) # Extract a specific position in the tuple cols = [(i[0:1][0]) for i in cols_args] attrs = [(i[1:]) for i in cols_args] else: # if not a list convert to list cols = val_to_list(cols_args) # Get col name from index cols = [c if is_str(c) else df.columns[c] for c in cols] # Check for missing columns if accepts_missing_cols is False: check_for_missing_columns(df, cols) # Filter by column data type filter_by_column_dtypes = val_to_list(filter_by_column_dtypes) if is_list_of_strings(filter_by_column_dtypes): # Get columns for every data type columns_filtered = filter_col_name_by_dtypes(df, filter_by_column_dtypes) # Intersect the columns filtered per datatype from the whole dataframe with the columns passed to the function cols = list(set(cols).intersection(columns_filtered)) # Return cols or cols an params if get_args is True: params = cols, attrs elif get_args is False: params = cols else: RaiseIfNot.value_error(get_args, ["True", "False"]) return params