fix

xorbitsai · Oct 16, 2023 · 40ee4f6 · 40ee4f6
1 parent f7153f7
commit 40ee4f6
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/xorbits/_mars/optimization/logical/tileable/column_pruning/column_pruning_rule.py b/python/xorbits/_mars/optimization/logical/tileable/column_pruning/column_pruning_rule.py
@@ -59,7 +59,11 @@ def _get_successor_required_columns(self, data: TileableData) -> Set[Any]:
  *[self._context[successor][data] for successor in successors]
  )
  # When getting the required columns of a DataFrameIndex node, we need to consider itself.
- if isinstance(data.op, DataFrameIndex) and len(data.dtypes) > 0:
+ if (
+ isinstance(data, BaseDataFrameData)
+ and isinstance(data.op, DataFrameIndex)
+ and len(data.dtypes) > 0
+ ):
  res = res.union(set(data.dtypes.index))
  return res
  else: