Optimize tile of DataFrame.__setitem__ by reducing time of generating chunk meta (#3140)

Author: Xuye (Chris) Qin

benchmarks/asv_bench/benchmarks/graph_builder.py

@@ -30,3 +30,10 @@ def setup(self):
     def time_filter(self):
         df = self.df[self.df["a"] < 0.8]
         build_graph([df], tile=True)
+
+    def time_setitem(self):
+        df2 = self.df.copy()
+        df2["k"] = df2["c"]
+        df2["l"] = df2["a"] * (1 - df2["d"])
+        df2["m"] = df2["e"] * (1 + df2["d"]) * (1 - df2["h"])
+        build_graph([df2], tile=True)

benchmarks/tpch/run_queries.py

@@ -1051,8 +1051,12 @@ def main():
     queries = (
         set(x.lower().strip() for x in args.query.split(",")) if args.query else None
     )
-    mars.new_session(endpoint)
-    run_queries(folder, use_arrow_dtype=use_arrow_dtype)
+    sess = mars.new_session(endpoint)
+    try:
+        run_queries(folder, use_arrow_dtype=use_arrow_dtype)
+    finally:
+        if endpoint is None:
+            sess.stop_server()
 
 
 if __name__ == "__main__":

mars/dataframe/align.py

@@ -39,6 +39,7 @@
     build_split_idx_to_origin_idx,
     filter_index_value,
     hash_index,
+    is_index_value_identical,
 )
 
 
@@ -974,6 +975,10 @@ def align_dataframe_series(left, right, axis="columns"):
 
 
 def align_series_series(left, right):
+    if is_index_value_identical(left, right):
+        # index identical, skip align
+        return left.nsplits, left.chunk_shape, left.chunks, right.chunks
+
     left_index_chunks = [c.index_value for c in left.chunks]
     right_index_chunks = [c.index_value for c in right.chunks]
 
@@ -988,9 +993,6 @@ def align_series_series(left, right):
 
     left_chunks = _gen_series_chunks(splits, out_chunk_shape, 0, left)
     right_chunks = _gen_series_chunks(splits, out_chunk_shape, 1, right)
-    if _is_index_identical(left_index_chunks, right_index_chunks):
-        index_nsplits = left.nsplits[0]
-    else:
-        index_nsplits = [np.nan for _ in range(out_chunk_shape[0])]
+    index_nsplits = [np.nan for _ in range(out_chunk_shape[0])]
     nsplits = [index_nsplits]
     return nsplits, out_chunk_shape, left_chunks, right_chunks

mars/dataframe/indexing/setitem.py

@@ -24,7 +24,7 @@
 from ..core import DATAFRAME_TYPE, SERIES_TYPE, DataFrame
 from ..initializer import DataFrame as asframe, Series as asseries
 from ..operands import DataFrameOperand, DataFrameOperandMixin
-from ..utils import parse_index
+from ..utils import parse_index, is_index_value_identical
 
 # in pandas 1.0.x, __setitem__ with a list with missing items are not allowed
 _allow_set_missing_list = pd_release_version[:2] >= (1, 1)
@@ -161,16 +161,7 @@ def tile(cls, op: "DataFrameSetitem"):
             rechunk_arg = {}
 
             # check if all chunk's index_value are identical
-            target_chunk_index_values = [
-                c.index_value for c in target.chunks if c.index[1] == 0
-            ]
-            value_chunk_index_values = [v.index_value for v in value.chunks]
-            is_identical = len(target_chunk_index_values) == len(
-                value_chunk_index_values
-            ) and all(
-                c.key == v.key
-                for c, v in zip(target_chunk_index_values, value_chunk_index_values)
-            )
+            is_identical = is_index_value_identical(target, value)
             if not is_identical:
                 # do rechunk
                 if any(np.isnan(s) for s in target.nsplits[0]) or any(
@@ -202,8 +193,8 @@ def tile(cls, op: "DataFrameSetitem"):
 
         out_chunks = []
         nsplits = [list(ns) for ns in target.nsplits]
-
         nsplits[1][-1] += len(append_cols)
+        nsplits = tuple(tuple(ns) for ns in nsplits)
 
         column_chunk_shape = target.chunk_shape[1]
         for c in target.chunks:
@@ -239,26 +230,27 @@ def tile(cls, op: "DataFrameSetitem"):
 
                     chunk_inputs = [c, value_chunk]
 
-                dtypes, shape, columns_value = c.dtypes, c.shape, c.columns_value
-
+                shape = c.shape
                 if append_cols and c.index[-1] == column_chunk_shape - 1:
                     # some columns appended at the last column of chunks
                     shape = (shape[0], shape[1] + len(append_cols))
-                    dtypes = pd.concat([dtypes, out.dtypes.iloc[-len(append_cols) :]])
-                    columns_value = parse_index(dtypes.index, store_data=True)
 
                 result_chunk = chunk_op.new_chunk(
                     chunk_inputs,
                     shape=shape,
-                    dtypes=dtypes,
-                    index_value=c.index_value,
-                    columns_value=columns_value,
                     index=c.index,
                 )
+                result_chunk._set_tileable_meta(
+                    tileable_key=out.key,
+                    nsplits=nsplits,
+                    index_value=out.index_value,
+                    columns_value=out.columns_value,
+                    dtypes=out.dtypes,
+                )
             out_chunks.append(result_chunk)
 
         params = out.params
-        params["nsplits"] = tuple(tuple(ns) for ns in nsplits)
+        params["nsplits"] = nsplits
         params["chunks"] = out_chunks
         new_op = op.copy()
         return new_op.new_tileables(op.inputs, kws=[params])
@@ -270,10 +262,17 @@ def estimate_size(cls, ctx: dict, op: "DataFrameSetitem"):
 
     @classmethod
     def execute(cls, ctx, op: "DataFrameSetitem"):
-        target = ctx[op.target.key].copy()
+        target = ctx[op.target.key]
+        # only deep copy when updating
+        indexes = (
+            (op.indexes,)
+            if not isinstance(op.indexes, (tuple, list, set))
+            else op.indexes
+        )
+        deep = bool(set(indexes) & set(target.columns))
+        target = ctx[op.target.key].copy(deep=deep)
         value = ctx[op.value.key] if not np.isscalar(op.value) else op.value
         try:
-
             target[op.indexes] = value
         except KeyError:
             if _allow_set_missing_list:  # pragma: no cover

mars/dataframe/utils.py

@@ -732,6 +732,28 @@ def build_concatenated_rows_frame(df):
     )
 
 
+def is_index_value_identical(left: TileableType, right: TileableType) -> bool:
+    if (
+        left.index_value.key == right.index_value.key
+        and not np.isnan(sum(left.nsplits[0]))
+        and not np.isnan(sum(right.nsplits[0]))
+        and left.nsplits[0] == right.nsplits[0]
+    ):
+        is_identical = True
+    else:
+        target_chunk_index_values = [
+            c.index_value for c in left.chunks if len(c.index) <= 1 or c.index[1] == 0
+        ]
+        value_chunk_index_values = [v.index_value for v in right.chunks]
+        is_identical = len(target_chunk_index_values) == len(
+            value_chunk_index_values
+        ) and all(
+            c.key == v.key
+            for c, v in zip(target_chunk_index_values, value_chunk_index_values)
+        )
+    return is_identical
+
+
 def _filter_range_index(pd_range_index, min_val, min_val_close, max_val, max_val_close):
     if is_pd_range_empty(pd_range_index):
         return pd_range_index