diff --git a/tests/connectors/test_excel.py b/tests/connectors/test_excel.py index 41fbc458..d2513e7b 100644 --- a/tests/connectors/test_excel.py +++ b/tests/connectors/test_excel.py @@ -1,4 +1,5 @@ import wrangles +import pandas as pd from wrangles.connectors import memory @@ -30,4 +31,209 @@ def test_default_write(): assert ( data["columns"] == ["header1", "header2"] and len(data["data"]) == 5 - ) \ No newline at end of file + ) + + +def test_recipe_wrangle_in_batch_writes_all_rows_to_excel_sheet(): + """ + Test the WranglesXL output connector path when a recipe wrangle is used + inside a batch. The Excel sheet output should receive the full combined + dataframe, not only one batch. + """ + memory.clear() + wrangles.recipe.run( + """ + read: + - test: + rows: 1000 + values: + header1: value1 + wrangles: + - batch: + batch_size: 100 + wrangles: + - recipe: + wrangles: + - convert.case: + input: header1 + case: upper + write: + - excel.sheet: + name: Partial + write: + - excel.sheet: + name: Final + """ + ) + + excel_outputs = [ + v + for v in memory.dataframes.values() + if v.get("connector") == "excel.sheet.write" + ] + memory.clear() + + assert len(excel_outputs) == 1 + assert excel_outputs[0]["name"] == "Final" + assert len(excel_outputs[0]["data"]) == 1000 + + +def test_excel_sheet_append_accumulates_repeated_writes(): + """ + WranglesXL may receive repeated writes to the same sheet when work is + batched. Default append behavior should accumulate rows in one payload. + """ + memory.clear() + df = pd.DataFrame({"header1": ["value1"] * 1000}) + for start in range(0, 1000, 100): + wrangles.connectors.excel.sheet.write( + df.iloc[start:start + 100], + name="Results" + ) + + excel_outputs = [ + v + for v in memory.dataframes.values() + if v.get("connector") == "excel.sheet.write" + ] + memory.clear() + + assert len(excel_outputs) == 1 + assert excel_outputs[0]["name"] == "Results" + assert len(excel_outputs[0]["data"]) == 1000 + + +def test_excel_sheet_append_aligns_dynamic_batch_columns_by_name(): + """ + Dynamic dictionary keys can create different columns in each batch. + Accumulated Excel output must union the columns and align values by name + instead of stacking each batch positionally. + """ + memory.clear() + batches = [ + pd.DataFrame({"ID": [1], "A": [1], "X": [97]}), + pd.DataFrame({"ID": [2], "B": [2], "Y": [98]}), + pd.DataFrame({"ID": [3], "C": [3], "Z": [99]}), + pd.DataFrame({"ID": [4], "D": [4], "Zz": [100]}), + ] + + for df in batches: + wrangles.connectors.excel.sheet.write(df, name="Results") + + excel_outputs = [ + v + for v in memory.dataframes.values() + if v.get("connector") == "excel.sheet.write" + ] + memory.clear() + + assert len(excel_outputs) == 1 + assert excel_outputs[0]["columns"] == [ + "ID", "A", "X", "B", "Y", "C", "Z", "D", "Zz" + ] + assert excel_outputs[0]["data"] == [ + [1, 1, 97, "", "", "", "", "", ""], + [2, "", "", 2, 98, "", "", "", ""], + [3, "", "", "", "", 3, 99, "", ""], + [4, "", "", "", "", "", "", 4, 100], + ] + + +def test_excel_sheet_overwrite_accumulates_repeated_writes(): + """ + Batched WranglesXL runs may emit repeated overwrite writes to the same + sheet. The connector should still return one full payload so Excel replaces + the sheet with all rows, not just the final batch. + """ + memory.clear() + df = pd.DataFrame({"header1": ["value1"] * 1000}) + for start in range(0, 1000, 100): + wrangles.connectors.excel.sheet.write( + df.iloc[start:start + 100], + name="Results", + action="overwrite" + ) + + excel_outputs = [ + v + for v in memory.dataframes.values() + if v.get("connector") == "excel.sheet.write" + ] + memory.clear() + + assert len(excel_outputs) == 1 + assert excel_outputs[0]["name"] == "Results" + assert excel_outputs[0]["action"] == "overwrite" + assert len(excel_outputs[0]["data"]) == 1000 + + +def test_excel_sheet_overwrite_aligns_reordered_columns_by_name(): + """ + Overwrite batches with the same columns in a different order must retain + the first payload's column order without shifting values. + """ + memory.clear() + wrangles.connectors.excel.sheet.write( + pd.DataFrame({"ID": [1], "A": [1], "X": [97]}), + name="Results", + action="overwrite" + ) + wrangles.connectors.excel.sheet.write( + pd.DataFrame({"X": [98], "ID": [2], "A": [2]}), + name="Results", + action="overwrite" + ) + + excel_outputs = [ + v + for v in memory.dataframes.values() + if v.get("connector") == "excel.sheet.write" + ] + memory.clear() + + assert len(excel_outputs) == 1 + assert excel_outputs[0]["columns"] == ["ID", "A", "X"] + assert excel_outputs[0]["data"] == [ + [1, 1, 97], + [2, 2, 98], + ] + assert excel_outputs[0]["action"] == "overwrite" + + +def test_excel_sheet_overwrite_uses_append_after_first_external_batch(): + """ + WranglesXL can execute each batch as a separate Python run. In that case, + in-memory accumulation is not available, so later overwrite batches must be + returned as append actions. + """ + df = pd.DataFrame({"header1": ["value1"] * 100}) + + memory.clear() + wrangles.connectors.excel.sheet.write( + df, + name="Results", + action="overwrite", + variables={"batch_number": 1, "batch_total": 10} + ) + first_batch = [ + v + for v in memory.dataframes.values() + if v.get("connector") == "excel.sheet.write" + ][0] + + memory.clear() + wrangles.connectors.excel.sheet.write( + df, + name="Results", + action="overwrite", + variables={"batch_number": 2, "batch_total": 10} + ) + second_batch = [ + v + for v in memory.dataframes.values() + if v.get("connector") == "excel.sheet.write" + ][0] + memory.clear() + + assert first_batch["action"] == "overwrite" + assert second_batch["action"] == "append" diff --git a/tests/connectors/test_recipe.py b/tests/connectors/test_recipe.py index f7fc3352..10799612 100644 --- a/tests/connectors/test_recipe.py +++ b/tests/connectors/test_recipe.py @@ -137,6 +137,42 @@ def test_wrangles(): ) assert df["header1"][0] == "VALUE1" + +def test_recipe_wrangle_in_batch_does_not_write_partial_batches(): + """ + Test that a recipe used as a wrangle does not run its own write step + once per batch. The outer recipe should write the final combined result. + """ + memory.clear() + + wrangles.recipe.run( + """ + read: + - test: + rows: 1000 + values: + header1: value1 + wrangles: + - batch: + batch_size: 100 + wrangles: + - recipe: + wrangles: + - convert.case: + input: header1 + case: upper + write: + - memory: + id: partial_recipe_write + write: + - memory: + id: final_recipe_write + """ + ) + + assert "partial_recipe_write" not in memory.dataframes + assert len(memory.dataframes["final_recipe_write"]["data"]) == 1000 + def test_write(): """ Test write defined within the recipe diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index 82ed87df..c3c8b08a 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -5715,6 +5715,41 @@ def record_batch(df): number_of_batches[0] == 1000 ) + def test_batch_dynamic_dictionary_columns_align_by_name(self): + """ + Dynamic dictionary keys can produce different columns in every batch. + The combined result must union those columns and align values by name. + """ + df = wrangles.recipe.run( + """ + wrangles: + - batch: + batch_size: 1 + wrangles: + - split.dictionary: + input: Spec Dicts + """, + dataframe=pd.DataFrame({ + "ID": [1, 2, 3, 4], + "Spec Dicts": [ + {"A": "1", "X": "97"}, + {"B": "2", "Y": "98"}, + {"C": "3", "Z": "99"}, + {"D": "4", "Zz": "100"}, + ] + }) + ) + + assert df.columns.tolist() == [ + "ID", "Spec Dicts", "A", "X", "B", "Y", "C", "Z", "D", "Zz" + ] + assert df[["A", "X", "B", "Y", "C", "Z", "D", "Zz"]].values.tolist() == [ + ["1", "97", "", "", "", "", "", ""], + ["", "", "2", "98", "", "", "", ""], + ["", "", "", "", "3", "99", "", ""], + ["", "", "", "", "", "", "4", "100"], + ] + def test_batch_preserves_order(self): """ Test batch preserves diff --git a/wrangles/connectors/excel.py b/wrangles/connectors/excel.py index 9e7102b4..5bf7b004 100644 --- a/wrangles/connectors/excel.py +++ b/wrangles/connectors/excel.py @@ -5,12 +5,98 @@ from . import memory as _memory import logging as _logging + +def _append_rows_by_column(saved: dict, df: _pd.DataFrame) -> bool: + """ + Append a dataframe to an orient="split" payload, aligning values by column + name and adding newly encountered columns in first-seen order. + """ + new_data = df.to_dict(orient="split") + saved_columns = saved["columns"] + new_columns = new_data["columns"] + + # Duplicate labels cannot be aligned by name unambiguously. Preserve the + # existing behavior for identical layouts, but leave different layouts as + # separate writes rather than risking a positional shift. + if ( + saved_columns != new_columns + and ( + not _pd.Index(saved_columns).is_unique + or not _pd.Index(new_columns).is_unique + ) + ): + return False + + combined_columns = saved_columns + [ + column + for column in new_columns + if column not in saved_columns + ] + + if combined_columns == saved_columns == new_columns: + saved["data"].extend(new_data["data"]) + else: + added_columns = len(combined_columns) - len(saved_columns) + if added_columns: + saved["data"] = [ + list(row) + [""] * added_columns + for row in saved["data"] + ] + + column_positions = { + column: position + for position, column in enumerate(combined_columns) + } + for row in new_data["data"]: + aligned_row = [""] * len(combined_columns) + for column, value in zip(new_columns, row): + aligned_row[column_positions[column]] = value + saved["data"].append(aligned_row) + + saved["columns"] = combined_columns + + saved["index"].extend(new_data["index"]) + return True + + class sheet(): _schema = {} - def write(df: _pd.DataFrame, **kwargs): + def write(df: _pd.DataFrame, variables: dict = None, **kwargs): _logging.info(f": Saving data for Excel Sheet") + if variables is None: + variables = {} + + action = kwargs.get("action", "append") + try: + batch_number = int(variables.get("batch_number", 1)) + batch_total = int(variables.get("batch_total", 1)) + except (TypeError, ValueError): + batch_number = 1 + batch_total = 1 + + if action == "overwrite" and batch_total > 1 and batch_number > 1: + kwargs["action"] = "append" + action = "append" + + name = kwargs.get("name") + cell = kwargs.get("cell") + + if action in ("append", "overwrite"): + for saved in reversed(list(_memory.dataframes.values())): + if ( + isinstance(saved, dict) + and saved.get("connector") == "excel.sheet.write" + and saved.get("name") == name + and saved.get("cell") == cell + and saved.get("action", "append") in ("append", "overwrite") + ): + if _append_rows_by_column(saved, df): + if action == "overwrite": + saved["action"] = "overwrite" + return + _memory.write( df, connector = "excel.sheet.write", diff --git a/wrangles/recipe_wrangles/main.py b/wrangles/recipe_wrangles/main.py index 83426351..4ff89680 100644 --- a/wrangles/recipe_wrangles/main.py +++ b/wrangles/recipe_wrangles/main.py @@ -1417,7 +1417,7 @@ def recipe( input: _Union[str, int, list] = None, output: _Union[str, list] = None, name: str = None, - variables = {}, + variables = None, functions: _Union[_types.FunctionType, list] = [], **kwargs ) -> _pd.DataFrame: @@ -1437,6 +1437,8 @@ def recipe( type: object description: A dictionary of variables to pass to the recipe """ + if variables is None: + variables = {} if not name: name = kwargs df_temp = df.copy() # copy of the original df @@ -1449,18 +1451,25 @@ def recipe( if output is None and input is not None: output = input + recipe_object, functions = _recipe._load_recipe( + name, + variables=variables, + functions=functions + ) + recipe_object.pop('write', None) + # If output columns are specified, only apply to those if output: if not isinstance(output, list): output = [output] df[output] = _recipe.run( - name, + recipe_object, variables=variables, functions=functions, dataframe=df_temp )[output] else: df = _recipe.run( - name, + recipe_object, variables=variables, functions=functions, dataframe=df_temp