fix

python/xorbits/_mars/optimization/logical/tileable/column_pruning/input_column_selector.py

@@ -164,9 +164,10 @@ def df_groupby_agg_select_function(
  ret = {}
  # group by a series
  groupby_series = False
- if isinstance(by, list) and len(by) == 1 and isinstance(by[0], BaseSeriesData):
+ if isinstance(by, list) and all([isinstance(_by, BaseSeriesData) for _by in by]):
  groupby_series = True
- ret[by[0]] = {by[0].name}
+ for _by in by:
+ ret[_by] = {_by.name}
 
  if isinstance(inp, BaseSeriesData):
  ret[inp] = {inp.name}

python/xorbits/_mars/optimization/logical/tileable/column_pruning/tests/test_input_column_selector.py

@@ -116,6 +116,18 @@ def test_df_groupby_agg():
  assert labels.data in input_columns
  assert input_columns[labels.data] == {"label"}
 
+ label1 = Series([1, 1, 1, 1], name="label1")
+ label2 = Series([2, 2, 3, 3], name="label2")
+ s = df.groupby(by=[label1, label2]).sum()
+ input_columns = InputColumnSelector.select(s.data, {"foo"})
+ assert len(input_columns) == 3
+ assert df.data in input_columns
+ assert input_columns[df.data] == {"foo"}
+ assert label1.data in input_columns
+ assert input_columns[label1.data] == {"label1"}
+ assert label2.data in input_columns
+ assert input_columns[label2.data] == {"label2"}
+
 
 @pytest.mark.skip(reason="group by index is not supported yet")
 def test_df_groupby_index_agg():