From 87a17aed476f905a30414408b8fa2f6125397682 Mon Sep 17 00:00:00 2001
From: David Nowinsky <david.nowinsky@toucantoco.com>
Date: Wed, 17 Apr 2024 10:49:55 +0200
Subject: [PATCH] fix(server): correctly generate one parent and one child row
 even if labels are the same TCTC-8331  (#2072)

* test(server): add row when parent and child have the same label

In mongo, it generates one child and one parent row.
In pandas, it generates 2 parent rows (bug).

* fix(server): correctly generate one parent and one child dow even if labels are the same (pandas, waterfall step)

* style(server): lint

* docs(server): changelog note
---
 server/CHANGELOG.md                               |  4 ++++
 .../backends/pandas_executor/steps/waterfall.py   | 15 ++++++---------
 .../fixtures/waterfall/with_backfill.yaml         | 15 ++++++++++++++-
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/server/CHANGELOG.md b/server/CHANGELOG.md
index 222e48a5d6..317055386e 100644
--- a/server/CHANGELOG.md
+++ b/server/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Pandas: Waterfall step: generate one parent and one child row even if labels are the same (instead of two parent rows)
+
 ## [0.42.2] - 2024-03-12
 
 ### Fixed
diff --git a/server/src/weaverbird/backends/pandas_executor/steps/waterfall.py b/server/src/weaverbird/backends/pandas_executor/steps/waterfall.py
index 4d7a07fd0c..830f35611b 100644
--- a/server/src/weaverbird/backends/pandas_executor/steps/waterfall.py
+++ b/server/src/weaverbird/backends/pandas_executor/steps/waterfall.py
@@ -127,7 +127,11 @@ def _merge(step: WaterfallStep, start_df: DataFrame, end_df: DataFrame) -> DataF
             {RESULT_COLUMN: "sum"}
         )
         parents_results[step.labelsColumn] = parents_results[step.parentsColumn]
+        parents_results[TYPE_WATERFALL_COLUMN] = "parent"
+        merged_df[TYPE_WATERFALL_COLUMN] = "child"
         return pd.concat([merged_df, parents_results])
+    else:
+        merged_df[TYPE_WATERFALL_COLUMN] = "parent"
     return merged_df
 
 
@@ -207,20 +211,13 @@ def _compute_parents_and_children(step: WaterfallStep, merged_df: DataFrame) ->
     """Computes children values (the deepest nesting level)"""
     result_df = DataFrame({LABEL_WATERFALL_COLUMN: merged_df[step.labelsColumn].astype(str)})
     result_df[step.groupby] = merged_df[step.groupby]
+    result_df[TYPE_WATERFALL_COLUMN] = merged_df[TYPE_WATERFALL_COLUMN]
 
     if step.parentsColumn is None:
-        result_df[TYPE_WATERFALL_COLUMN] = "parent"
         result_df[_VQB_SORT_ORDER] = 2
     else:
         result_df[GROUP_WATERFALL_COLUMN] = merged_df[step.parentsColumn]
-        result_df[TYPE_WATERFALL_COLUMN] = numpy.where(
-            result_df[LABEL_WATERFALL_COLUMN] == (result_df[GROUP_WATERFALL_COLUMN]),
-            "parent",
-            "child",
-        )
-        result_df[_VQB_SORT_ORDER] = numpy.where(
-            result_df[LABEL_WATERFALL_COLUMN] == (result_df[GROUP_WATERFALL_COLUMN]), 2, 1
-        )
+        result_df[_VQB_SORT_ORDER] = numpy.where(result_df[TYPE_WATERFALL_COLUMN] == "parent", 2, 1)
 
     result_df[step.valueColumn] = merged_df[RESULT_COLUMN]
     return result_df
diff --git a/server/tests/backends/fixtures/waterfall/with_backfill.yaml b/server/tests/backends/fixtures/waterfall/with_backfill.yaml
index 46be40d240..f0ea7113dc 100644
--- a/server/tests/backends/fixtures/waterfall/with_backfill.yaml
+++ b/server/tests/backends/fixtures/waterfall/with_backfill.yaml
@@ -259,6 +259,11 @@ input:
     country: Germany
     month: 01 April 2023
     price: 65076
+  # Only element with the same name as the group
+  - account: England
+    country: England
+    month: 01 January 2023
+    price: 123456
   schema:
     fields:
     - name: account
@@ -275,7 +280,7 @@ expected:
   - GROUP_waterfall: 01 January 2023
     LABEL_waterfall: 01 January 2023
     TYPE_waterfall: null
-    price: 119907
+    price: 243363
   - GROUP_waterfall: Germany
     LABEL_waterfall: Fermentum Industries
     TYPE_waterfall: child
@@ -296,6 +301,10 @@ expected:
     LABEL_waterfall: Dignissim Maecenas Foundation
     TYPE_waterfall: child
     price: -62202
+  - GROUP_waterfall: England
+    LABEL_waterfall: England
+    TYPE_waterfall: child
+    price: -123456
   - GROUP_waterfall: Germany
     LABEL_waterfall: Germany
     TYPE_waterfall: parent
@@ -304,6 +313,10 @@ expected:
     LABEL_waterfall: France
     TYPE_waterfall: parent
     price: -78287
+  - GROUP_waterfall: England
+    LABEL_waterfall: England
+    TYPE_waterfall: parent
+    price: -123456
   - GROUP_waterfall: 01 December 2023
     LABEL_waterfall: 01 December 2023
     TYPE_waterfall: null