Repository: mars-project/mars

Issue:
Support `DataFrame.cartesian_chunk`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**

`DataFrame.cartesian_chunk` is an extended API that is useful when users want to perform cartesian join computation. A function can be provided to perform calculation on two pandas DataFrames.


Discussion:

Resolving patch:
diff --git a/mars/dataframe/base/__init__.py b/mars/dataframe/base/__init__.py
--- a/mars/dataframe/base/__init__.py
+++ b/mars/dataframe/base/__init__.py
@@ -33,6 +33,7 @@
 from .memory_usage import df_memory_usage, series_memory_usage, index_memory_usage
 from .select_dtypes import select_dtypes
 from .map_chunk import map_chunk
+from .cartesian_chunk import cartesian_chunk
 from .rebalance import rebalance
 from .stack import stack
 from .explode import df_explode, series_explode
@@ -65,6 +66,7 @@ def _install():
         setattr(t, 'memory_usage', df_memory_usage)
         setattr(t, 'select_dtypes', select_dtypes)
         setattr(t, 'map_chunk', map_chunk)
+        setattr(t, 'cartesian_chunk', cartesian_chunk)
         setattr(t, 'rebalance', rebalance)
         setattr(t, 'stack', stack)
         setattr(t, 'explode', df_explode)
@@ -87,6 +89,7 @@ def _install():
         setattr(t, 'drop_duplicates', series_drop_duplicates)
         setattr(t, 'memory_usage', series_memory_usage)
         setattr(t, 'map_chunk', map_chunk)
+        setattr(t, 'cartesian_chunk', cartesian_chunk)
         setattr(t, 'rebalance', rebalance)
         setattr(t, 'explode', series_explode)
 
diff --git a/mars/dataframe/base/cartesian_chunk.py b/mars/dataframe/base/cartesian_chunk.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/base/cartesian_chunk.py
@@ -0,0 +1,207 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...custom_log import redirect_custom_log
+from ...serialize import KeyField, FunctionField, TupleField, \
+    DictField, StringField
+from ...tiles import TilesError
+from ...utils import check_chunks_unknown_shape, enter_current_session
+from ..operands import DataFrameOperand, DataFrameOperandMixin, OutputType
+from ..utils import build_df, build_series, build_empty_df, parse_index, \
+    validate_output_types, quiet_stdio
+
+
+class DataFrameCartesianChunk(DataFrameOperand, DataFrameOperandMixin):
+    _op_type_ = opcodes.CARTESIAN_CHUNK
+
+    _left = KeyField('left')
+    _right = KeyField('right')
+    _func = FunctionField('func')
+    _args = TupleField('args')
+    _kwargs = DictField('kwargs')
+    # for chunk
+    _tileable_op_key = StringField('tileable_op_key')
+
+    def __init__(self, left=None, right=None, func=None, args=None, kwargs=None,
+                 memory_scale=None, output_types=None, tileable_op_key=None, **kw):
+        super().__init__(_left=left, _right=right, _func=func, _args=args,
+                         _kwargs=kwargs, _memory_scale=memory_scale,
+                         _output_types=output_types,
+                         _tileable_op_key=tileable_op_key, **kw)
+
+    @property
+    def left(self):
+        return self._left
+
+    @property
+    def right(self):
+        return self._right
+
+    @property
+    def func(self):
+        return self._func
+
+    @property
+    def args(self):
+        return self._args
+
+    @property
+    def kwargs(self):
+        return self._kwargs
+
+    @property
+    def memory_scale(self):
+        return self._memory_scale or 2
+
+    @property
+    def tileable_op_key(self):
+        return self._tileable_op_key
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        self._left = self._inputs[0]
+        self._right = self._inputs[1]
+
+    @staticmethod
+    def _build_test_obj(obj):
+        return build_df(obj, size=2) if obj.ndim == 2 else \
+            build_series(obj, size=2, name=obj.name)
+
+    def __call__(self, left, right, index=None, dtypes=None):
+        test_left = self._build_test_obj(left)
+        test_right = self._build_test_obj(right)
+        output_type = self._output_types[0] if self._output_types else None
+
+        # try run to infer meta
+        try:
+            with np.errstate(all='ignore'), quiet_stdio():
+                obj = self._func(test_left, test_right, *self._args, **self._kwargs)
+        except:  # noqa: E722  # nosec  # pylint: disable=bare-except
+            if output_type == OutputType.series:
+                obj = pd.Series([], dtype=np.dtype(object))
+            elif output_type == OutputType.dataframe and dtypes is not None:
+                obj = build_empty_df(dtypes)
+            else:
+                raise TypeError('Cannot determine `output_type`, '
+                                'you have to specify it as `dataframe` or `series`, '
+                                'for dataframe, `dtypes` is required as well '
+                                'if output_type=\'dataframe\'')
+
+        if getattr(obj, 'ndim', 0) == 1 or output_type == OutputType.series:
+            shape = self._kwargs.pop('shape', (np.nan,))
+            if index is None:
+                index = obj.index
+            index_value = parse_index(index, left, right,
+                                      self._func, self._args, self._kwargs)
+            return self.new_series([left, right], dtype=obj.dtype,
+                                   shape=shape, index_value=index_value,
+                                   name=obj.name)
+        else:
+            dtypes = dtypes if dtypes is not None else obj.dtypes
+            # dataframe
+            shape = (np.nan, len(dtypes))
+            columns_value = parse_index(dtypes.index, store_data=True)
+            if index is None:
+                index = obj.index
+            index_value = parse_index(index, left, right,
+                                      self._func, self._args, self._kwargs)
+            return self.new_dataframe([left, right], shape=shape,
+                                      dtypes=dtypes, index_value=index_value,
+                                      columns_value=columns_value)
+
+    @classmethod
+    def tile(cls, op: "DataFrameCartesianChunk"):
+        left = op.left
+        right = op.right
+        out = op.outputs[0]
+
+        if left.ndim == 2 and left.chunk_shape[1] > 1:
+            check_chunks_unknown_shape([left], TilesError)
+            # if left is a DataFrame, make sure 1 chunk on axis columns
+            left = left.rechunk({1: left.shape[1]})._inplace_tile()
+        if right.ndim == 2 and right.chunk_shape[1] > 1:
+            check_chunks_unknown_shape([right], TilesError)
+            # if right is a DataFrame, make sure 1 chunk on axis columns
+            right = right.rechunk({1: right.shape[1]})._inplace_tile()
+
+        out_chunks = []
+        nsplits = [[]] if out.ndim == 1 else [[], [out.shape[1]]]
+        i = 0
+        for left_chunk in left.chunks:
+            for right_chunk in right.chunks:
+                chunk_op = op.copy().reset_key()
+                chunk_op._tileable_op_key = op.key
+                if op.output_types[0] == OutputType.dataframe:
+                    shape = (np.nan, out.shape[1])
+                    index_value = parse_index(out.index_value.to_pandas(),
+                                              left_chunk, right_chunk,
+                                              op.func, op.args, op.kwargs)
+                    out_chunk = chunk_op.new_chunk([left_chunk, right_chunk],
+                                                   shape=shape,
+                                                   index_value=index_value,
+                                                   columns_value=out.columns_value,
+                                                   index=(i, 0))
+                    out_chunks.append(out_chunk)
+                    nsplits[0].append(out_chunk.shape[0])
+                else:
+                    shape = (np.nan,)
+                    index_value = parse_index(out.index_value.to_pandas(),
+                                              left_chunk, right_chunk,
+                                              op.func, op.args, op.kwargs)
+                    out_chunk = chunk_op.new_chunk([left_chunk, right_chunk],
+                                                   shape=shape,
+                                                   index_value=index_value,
+                                                   name=out.name,
+                                                   index=(i,))
+                    out_chunks.append(out_chunk)
+                    nsplits[0].append(out_chunk.shape[0])
+                i += 1
+
+        params = out.params
+        params['nsplits'] = tuple(tuple(ns) for ns in nsplits)
+        params['chunks'] = out_chunks
+        new_op = op.copy()
+        return new_op.new_tileables(op.inputs, kws=[params])
+
+    @classmethod
+    @redirect_custom_log
+    @enter_current_session
+    def execute(cls, ctx, op: "DataFrameCartesianChunk"):
+        left, right = ctx[op.left.key], ctx[op.right.key]
+        ctx[op.outputs[0].key] = op.func(
+            left, right, *op.args, **(op.kwargs or dict()))
+
+
+def cartesian_chunk(left, right, func, args=(), **kwargs):
+    output_type = kwargs.pop('output_type', None)
+    output_types = kwargs.pop('output_types', None)
+    object_type = kwargs.pop('object_type', None)
+    output_types = validate_output_types(
+        output_type=output_type, output_types=output_types, object_type=object_type)
+    output_type = output_types[0] if output_types else None
+    if output_type:
+        output_types = [output_type]
+    index = kwargs.pop('index', None)
+    dtypes = kwargs.pop('dtypes', None)
+    memory_scale = kwargs.pop('memory_scale', None)
+
+    op = DataFrameCartesianChunk(left=left, right=right, func=func,
+                                 args=args, kwargs=kwargs,
+                                 output_types=output_types,
+                                 memory_scale=memory_scale)
+    return op(left, right, index=index, dtypes=dtypes)
diff --git a/mars/dataframe/base/map_chunk.py b/mars/dataframe/base/map_chunk.py
--- a/mars/dataframe/base/map_chunk.py
+++ b/mars/dataframe/base/map_chunk.py
@@ -18,10 +18,11 @@
 from ... import opcodes
 from ...custom_log import redirect_custom_log
 from ...serialize import KeyField, FunctionField, TupleField, DictField, StringField
-from ...utils import enter_current_session
+from ...tiles import TilesError
+from ...utils import enter_current_session, check_chunks_unknown_shape
 from ..operands import DataFrameOperand, DataFrameOperandMixin, OutputType
 from ..utils import build_df, build_empty_df, build_series, \
-    parse_index, validate_output_types
+    parse_index, validate_output_types, quiet_stdio
 
 
 class DataFrameMapChunk(DataFrameOperand, DataFrameOperandMixin):
@@ -72,7 +73,7 @@ def __call__(self, df_or_series, index=None, dtypes=None):
 
         # try run to infer meta
         try:
-            with np.errstate(all='ignore'):
+            with np.errstate(all='ignore'), quiet_stdio():
                 obj = self._func(test_obj, *self._args, **self._kwargs)
         except:  # noqa: E722  # nosec
             if df_or_series.ndim == 1 or output_type == OutputType.series:
@@ -122,6 +123,7 @@ def tile(cls, op: "DataFrameMapChunk"):
         out = op.outputs[0]
 
         if inp.ndim == 2 and inp.chunk_shape[1] > 1:
+            check_chunks_unknown_shape([inp], TilesError)
             # if input is a DataFrame, make sure 1 chunk on axis columns
             inp = inp.rechunk({1: inp.shape[1]})._inplace_tile()
 


Test patch:
diff --git a/mars/dataframe/base/tests/test_base_execution.py b/mars/dataframe/base/tests/test_base_execution.py
--- a/mars/dataframe/base/tests/test_base_execution.py
+++ b/mars/dataframe/base/tests/test_base_execution.py
@@ -1462,6 +1462,85 @@ def f4(pdf):
         expected = raw2.apply(lambda x: x['a'], axis=1, result_type='expand')
         pd.testing.assert_frame_equal(result, expected)
 
+    def testCartesianChunkExecution(self):
+        rs = np.random.RandomState(0)
+        raw1 = pd.DataFrame({'a': rs.randint(3, size=10),
+                             'b': rs.rand(10)})
+        raw2 = pd.DataFrame({'c': rs.randint(3, size=10),
+                             'd': rs.rand(10),
+                             'e': rs.rand(10)})
+        df1 = from_pandas_df(raw1, chunk_size=(5, 1))
+        df2 = from_pandas_df(raw2, chunk_size=(5, 1))
+
+        def f(c1, c2):
+            c1, c2 = c1.copy(), c2.copy()
+            c1['x'] = 1
+            c2['x'] = 1
+            r = c1.merge(c2, on='x')
+            r = r[(r['b'] > r['d']) & (r['b'] < r['e'])]
+            return r[['a', 'c']]
+
+        rr = df1.cartesian_chunk(df2, f)
+
+        result = self.executor.execute_dataframe(rr, concat=True)[0]
+        expected = f(raw1, raw2)
+        pd.testing.assert_frame_equal(
+            result.sort_values(by=['a', 'c']).reset_index(drop=True),
+            expected.sort_values(by=['a', 'c']).reset_index(drop=True))
+
+        def f2(c1, c2):
+            r = f(c1, c2)
+            return r['a'] + r['c']
+
+        rr = df1.cartesian_chunk(df2, f2)
+
+        result = self.executor.execute_dataframe(rr, concat=True)[0]
+        expected = f2(raw1, raw2)
+        pd.testing.assert_series_equal(
+            result.sort_values().reset_index(drop=True),
+            expected.sort_values().reset_index(drop=True))
+
+        size_res = self.executor.execute_dataframe(rr, mock=True)[0][0]
+        self.assertGreater(size_res, 0)
+
+        def f3(c1, c2):
+            cr = pd.DataFrame()
+            cr['a'] = c1.str.slice(1).astype(np.int64)
+            cr['x'] = 1
+            cr2 = pd.DataFrame()
+            cr2['b'] = c2.str.slice(1).astype(np.int64)
+            cr2['x'] = 1
+            return cr.merge(cr2, on='x')[['a', 'b']]
+
+        s_raw = pd.Series([f's{i}' for i in range(10)])
+        series = from_pandas_series(s_raw, chunk_size=5)
+
+        rr = series.cartesian_chunk(series, f3, output_type='dataframe',
+                                    dtypes=pd.Series([np.dtype(np.int64)] * 2,
+                                                     index=['a', 'b']))
+
+        result = self.executor.execute_dataframe(rr, concat=True)[0]
+        expected = f3(s_raw, s_raw)
+        pd.testing.assert_frame_equal(
+            result.sort_values(by=['a', 'b']).reset_index(drop=True),
+            expected.sort_values(by=['a', 'b']).reset_index(drop=True))
+
+        with self.assertRaises(TypeError):
+            _ = series.cartesian_chunk(series, f3)
+
+        def f4(c1, c2):
+            r = f3(c1, c2)
+            return r['a'] + r['b']
+
+        rr = series.cartesian_chunk(series, f4, output_type='series',
+                                    dtypes=np.dtype(np.int64))
+
+        result = self.executor.execute_dataframe(rr, concat=True)[0]
+        expected = f4(s_raw, s_raw)
+        pd.testing.assert_series_equal(
+            result.sort_values().reset_index(drop=True),
+            expected.sort_values().reset_index(drop=True))
+
     def testRebalanceExecution(self):
         raw = pd.DataFrame(np.random.rand(10, 3), columns=list('abc'))
         df = from_pandas_df(raw)