fix group by not working for boolean columns

briefercloud · Feb 3, 2025 · b307cab · b307cab
1 parent 998fbba
commit b307cab
Show file tree

Hide file tree

Showing 2 changed files with 174 additions and 8 deletions.
diff --git a/apps/api/src/python/visualization-v2.test.ts b/apps/api/src/python/visualization-v2.test.ts
@@ -156,7 +156,8 @@ df = pd.DataFrame({
   'amount': [11, 12, 13, 21, 22, 33],
   'price': [1.1, 1.2, 1.3, 2.1, 2.2, 3.3],
   'datetime': pd.to_datetime(['2021-01-01', '2021-01-02', '2021-01-03', '2021-02-01', '2021-02-02', '2021-03-03']),
-  'fruit': ['apple', 'banana', 'pineapple', 'banana', 'apple', 'apple']
+  'fruit': ['apple', 'banana', 'pineapple', 'banana', 'apple', 'apple'],
+  'is_apple': [True, False, False, False, True, True]
 })`
 
   const fruitDFColumn: DataFrameColumn = {
@@ -179,9 +180,20 @@ df = pd.DataFrame({
     name: 'datetime',
   }
 
+  const isAppleDFColumn: DataFrameColumn = {
+    type: 'bool',
+    name: 'is_apple',
+  }
+
   const df: DataFrame = {
     name: 'df',
-    columns: [fruitDFColumn, amountDFColumn, priceDFColumn, datetimeDFColumn],
+    columns: [
+      fruitDFColumn,
+      amountDFColumn,
+      priceDFColumn,
+      datetimeDFColumn,
+      isAppleDFColumn,
+    ],
   }
 
   it('it should compute normalized values for hundredPercentStackedColumn', async () => {
@@ -675,4 +687,158 @@ df = pd.DataFrame({
       },
     })
   })
+
+  it('should work properly when group by is on a boolean column', async () => {
+    await (
+      await pythonRunner.runPython('workspaceId', 'sessionId', code, () => {}, {
+        storeHistory: true,
+      })
+    ).promise
+
+    const input: VisualizationV2BlockInput = {
+      dataframeName: 'df',
+      chartType: 'groupedColumn',
+      xAxis: datetimeDFColumn,
+      xAxisName: null,
+      xAxisSort: 'ascending',
+      xAxisGroupFunction: 'month',
+      yAxes: [
+        {
+          id: 'yAxis-1',
+          name: null,
+          series: [
+            {
+              id: 'series-1',
+              chartType: null,
+              column: amountDFColumn,
+              aggregateFunction: 'sum',
+              groupBy: isAppleDFColumn,
+              name: null,
+              color: null,
+              groups: null,
+            },
+          ],
+        },
+      ],
+      histogramFormat: 'count',
+      histogramBin: { type: 'auto' },
+      filters: [],
+      dataLabels: {
+        show: false,
+        frequency: 'all',
+      },
+    }
+
+    const result = await (
+      await createVisualizationV2(
+        'workspaceId',
+        'sessionId',
+        df,
+        input,
+        manager,
+        pythonRunner.runPython
+      )
+    ).promise
+
+    const janAppleTotal = 11
+    const janNotAppleTotal = 12 + 13
+    const febAppleTotal = 22
+    const febNotAppleTotal = 21
+    const marAppleTotal = 33
+
+    expect(result).toEqual({
+      success: true,
+      tooManyDataPoints: false,
+      filters: [],
+      data: {
+        tooltip: {
+          trigger: 'axis',
+        },
+        legend: {},
+        grid: {
+          containLabel: true,
+        },
+        dataset: [
+          {
+            dimensions: ['datetime', 'series-1'],
+            source: [
+              {
+                'series-1': janNotAppleTotal,
+                datetime: '2021-01-01 00:00:00',
+              },
+              {
+                'series-1': febNotAppleTotal,
+                datetime: '2021-02-01 00:00:00',
+              },
+            ],
+          },
+          {
+            dimensions: ['datetime', 'series-1'],
+            source: [
+              {
+                'series-1': janAppleTotal,
+                datetime: '2021-01-01 00:00:00',
+              },
+              {
+                'series-1': febAppleTotal,
+                datetime: '2021-02-01 00:00:00',
+              },
+              {
+                'series-1': marAppleTotal,
+                datetime: '2021-03-01 00:00:00',
+              },
+            ],
+          },
+        ],
+        xAxis: [
+          {
+            type: 'time',
+            axisPointer: {
+              type: 'shadow',
+            },
+            name: null,
+            nameLocation: 'middle',
+            max: 'dataMax',
+            min: 'dataMin',
+          },
+        ],
+        yAxis: [
+          {
+            type: 'value',
+            position: 'left',
+            name: null,
+            nameLocation: 'middle',
+          },
+        ],
+        series: [
+          {
+            id: 'series-1:False',
+            datasetIndex: 0,
+            yAxisIndex: 0,
+            name: 'False',
+            z: 0,
+            type: 'bar',
+            color: '#5470c6',
+            encode: {
+              x: 'datetime',
+              y: 'series-1',
+            },
+          },
+          {
+            id: 'series-1:True',
+            datasetIndex: 1,
+            yAxisIndex: 0,
+            name: 'True',
+            z: 0,
+            type: 'bar',
+            color: '#91cc75',
+            encode: {
+              x: 'datetime',
+              y: 'series-1',
+            },
+          },
+        ],
+      },
+    })
+  })
 })
diff --git a/apps/api/src/python/visualizations-v2.ts b/apps/api/src/python/visualizations-v2.ts
@@ -400,7 +400,7 @@ def _briefer_create_visualization(df, options):
                 for group in groups:
                     color_index += 1
                     dataset_index = len(data["dataset"])
-                    g_options = group_options.get(group) if group else {}
+                    g_options = group_options.get(group) if group is not None else {}
 
                     dimensions = [y_name]
                     if options["xAxis"]:
@@ -412,7 +412,7 @@ def _briefer_create_visualization(df, options):
                     }
 
                     for _, row in series_dataframe.iterrows():
-                        if group and row[series["groupBy"]["name"]] != group:
+                        if group is not None and row[series["groupBy"]["name"]] != group:
                             continue
 
                         y_value = row[y_name]
@@ -437,7 +437,7 @@ def _briefer_create_visualization(df, options):
                     data["dataset"].append(dataset)
 
                     id = f"{series['id']}"
-                    if group:
+                    if group is not None:
                         id = f"{id}:{group}"
 
                     serie = {
@@ -454,18 +454,18 @@ def _briefer_create_visualization(df, options):
                     if chart_type == "line":
                         serie["symbolSize"] = 1
 
-                    if not group:
+                    if group is None:
                         color = series.get("color") or colors[color_index % len(colors)]
                     elif g_options:
                         color = g_options.get("color") or colors[color_index % len(colors)]
                     else:
                         color = colors[color_index % len(colors)]
 
-                    if not group:
+                    if group is None:
                         serie["name"] = series.get("name") or series["column"]["name"]
                     elif group and g_options:
                         serie["name"] = g_options.get("name") or group
-                    elif group:
+                    elif group is not None:
                         serie["name"] = group
                     else:
                         serie["name"] = series["column"]["name"]