Repository: mars-project/mars

Issue:
[BUG] Error will be raised if `chunk_size` not specified for `md.read_sql_table`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

If `chunk_size` not specified for `md.read_sql_table`, error will be raised:

```
../../../utils.py:380: in _wrapped
    return func(*args, **kwargs)
../../../utils.py:468: in inner
    return func(*args, **kwargs)
../../../executor.py:641: in execute_tileable
    chunk_graph = chunk_graph_builder.build([tileable], tileable_graph=tileable_graph)
../../../utils.py:380: in _wrapped
    return func(*args, **kwargs)
../../../utils.py:468: in inner
    return func(*args, **kwargs)
../../../tiles.py:234: in build
    tiled = self._tile(tileable_data, tileable_graph)
../../../tiles.py:192: in _tile
    tds[0]._inplace_tile()
../../../core.py:161: in _inplace_tile
    return handler.inplace_tile(self)
../../../tiles.py:127: in inplace_tile
    dispatched = self.dispatch(to_tile.op)
../../../utils.py:380: in _wrapped
    return func(*args, **kwargs)
../../../tiles.py:115: in dispatch
    return op_cls.tile(op)
../read_sql_table.py:218: in tile
    row_chunk_sizes = decide_dataframe_chunk_sizes(df.shape, chunk_size, memory_usage)[0]
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

shape = (10, 1), chunk_size = {}, memory_usage = 670.0

    def decide_dataframe_chunk_sizes(shape, chunk_size, memory_usage):
        """
        Decide how a given DataFrame can be split into chunk.
    
        :param shape: DataFrame's shape
        :param chunk_size: if dict provided, it's dimension id to chunk size;
                           if provided, it's the chunk size for each dimension.
        :param memory_usage: pandas Series in which each column's memory usage
        :type memory_usage: pandas.Series
        :return: the calculated chunk size for each dimension
        :rtype: tuple
        """
        from ..config import options
    
        chunk_size = dictify_chunk_size(shape, chunk_size)
        average_memory_usage = memory_usage / shape[0]
    
        nleft = len(shape) - len(chunk_size)
        if nleft < 0:
            raise ValueError("chunks have more than two dimensions")
        if nleft == 0:
            return normalize_chunk_sizes(shape, tuple(chunk_size[j] for j in range(len(shape))))
    
        max_chunk_size = options.chunk_store_limit
    
        # for the row side, along axis 0
        if 0 not in chunk_size:
            row_chunk_size = []
            row_left_size = shape[0]
        else:
            row_chunk_size = normalize_chunk_sizes((shape[0],), (chunk_size[0],))[0]
            row_left_size = 0
        # for the column side, along axis 1
        if 1 not in chunk_size:
            col_chunk_size = []
            col_chunk_store = []
            col_left_size = shape[1]
        else:
            col_chunk_size = normalize_chunk_sizes((shape[1],), (chunk_size[1],))[0]
            acc = [0] + np.cumsum(col_chunk_size).tolist()
            col_chunk_store = [average_memory_usage[acc[i]: acc[i + 1]].sum()
                               for i in range(len(col_chunk_size))]
            col_left_size = 0
    
        while True:
            nbytes_occupied = np.prod([max(it) for it in (row_chunk_size, col_chunk_store) if it])
            dim_size = np.maximum(int(np.power(max_chunk_size / nbytes_occupied, 1 / float(nleft))), 1)
            # check col first
            if col_left_size > 0:
                cs = min(col_left_size, dim_size)
                col_chunk_size.append(cs)
                start = int(np.sum(col_chunk_size[:-1]))
>               col_chunk_store.append(average_memory_usage[start: start + cs].sum())
E               IndexError: invalid index to scalar variable.
```


Discussion:

Resolving patch:
diff --git a/mars/_version.py b/mars/_version.py
--- a/mars/_version.py
+++ b/mars/_version.py
@@ -15,7 +15,7 @@
 import subprocess
 import os
 
-version_info = (0, 4, 0, 'a1')
+version_info = (0, 4, 0, 'a2')
 _num_index = max(idx if isinstance(v, int) else 0
                  for idx, v in enumerate(version_info))
 __version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \
diff --git a/mars/dataframe/datasource/read_sql_table.py b/mars/dataframe/datasource/read_sql_table.py
--- a/mars/dataframe/datasource/read_sql_table.py
+++ b/mars/dataframe/datasource/read_sql_table.py
@@ -22,8 +22,9 @@
 from ...config import options
 from ...serialize import StringField, AnyField, BoolField, \
     ListField, Int64Field, Float64Field, BytesField
+from ...tensor.utils import normalize_chunk_sizes
 from ..operands import DataFrameOperand, DataFrameOperandMixin, ObjectType
-from ..utils import parse_index, decide_dataframe_chunk_sizes
+from ..utils import parse_index
 
 
 class DataFrameReadSQLTable(DataFrameOperand, DataFrameOperandMixin):
@@ -103,7 +104,7 @@ def _collect_info(self, engine_or_conn, table, columns, test_rows):
         from sqlalchemy import sql
 
         # fetch test DataFrame
-        query = sql.select(columns).limit(test_rows).select_from(table)
+        query = sql.select(columns).limit(test_rows)
         test_df = pd.read_sql(query, engine_or_conn, index_col=self._index_col,
                               coerce_float=self._coerce_float,
                               parse_dates=self._parse_dates)
@@ -213,9 +214,10 @@ def __call__(self, test_rows, chunk_size):
     def tile(cls, op):
         df = op.outputs[0]
 
-        memory_usage = op.row_memory_usage * df.shape[0]
         chunk_size = df.extra_params.raw_chunk_size or options.chunk_size
-        row_chunk_sizes = decide_dataframe_chunk_sizes(df.shape, chunk_size, memory_usage)[0]
+        if chunk_size is None:
+            chunk_size = (int(options.chunk_store_limit / op.row_memory_usage), df.shape[1])
+        row_chunk_sizes = normalize_chunk_sizes(df.shape, chunk_size)[0]
         offsets = np.cumsum((0,) + row_chunk_sizes)
 
         out_chunks = []
@@ -259,10 +261,20 @@ def execute(cls, ctx, op):
                     if icol not in column_names:
                         columns.append(table.columns[icol])
 
-            query = sa.sql.select(columns).limit(out.shape[0])
+            query = sa.sql.select(columns)
+            if len(table.primary_key) > 0:
+                # if table has primary key, sort as the order
+                query = query.order_by(*list(table.primary_key))
+            elif op.index_col:
+                # if no primary key, sort as the index_col
+                query = query.order_by(
+                    *[table.columns[col] for col in op.index_col])
+            else:
+                # at last, we sort by all the columns
+                query = query.order_by(*columns)
+            query = query.limit(out.shape[0])
             if op.offset > 0:
                 query = query.offset(op.offset)
-            query = query.select_from(table)
 
             df = pd.read_sql(query, engine, index_col=op.index_col,
                              coerce_float=op.coerce_float,


Test patch:
diff --git a/mars/dataframe/datasource/tests/test_datasource_execution.py b/mars/dataframe/datasource/tests/test_datasource_execution.py
--- a/mars/dataframe/datasource/tests/test_datasource_execution.py
+++ b/mars/dataframe/datasource/tests/test_datasource_execution.py
@@ -368,6 +368,7 @@ def testReadSQLTableExecution(self):
 
         with tempfile.TemporaryDirectory() as d:
             table_name = 'test'
+            table_name2 = 'test2'
             uri = 'sqlite:///' + os.path.join(d, 'test.db')
 
             test_df.to_sql(table_name, uri, index=False)
@@ -377,6 +378,7 @@ def testReadSQLTableExecution(self):
             pd.testing.assert_frame_equal(result, test_df)
 
             engine = sa.create_engine(uri)
+            m = sa.MetaData()
 
             try:
                 # test index_col and columns
@@ -388,7 +390,13 @@ def testReadSQLTableExecution(self):
                 del expected['c']
                 pd.testing.assert_frame_equal(result, expected)
 
-                table = sa.Table(table_name, sa.MetaData(), autoload=True,
+                # do not specify chunk_size
+                r = md.read_sql_table('test', engine.connect(),
+                                      index_col='a', columns=['b'])
+                result = self.executor.execute_dataframe(r, concat=True)[0]
+                pd.testing.assert_frame_equal(result, expected)
+
+                table = sa.Table(table_name, m, autoload=True,
                                  autoload_with=engine)
                 r = md.read_sql_table(table, engine, chunk_size=4,
                                       index_col=[table.columns['a'], table.columns['b']],
@@ -397,5 +405,20 @@ def testReadSQLTableExecution(self):
                 expected = test_df.copy(deep=True)
                 expected.set_index(['a', 'b'], inplace=True)
                 pd.testing.assert_frame_equal(result, expected)
+
+                # test primary key
+                sa.Table(table_name2, m,
+                         sa.Column('id', sa.Integer, primary_key=True),
+                         sa.Column('a', sa.Integer),
+                         sa.Column('b', sa.String),
+                         sa.Column('c', sa.Float))
+                m.create_all(engine)
+                test_df = test_df.copy(deep=True)
+                test_df.index.name = 'id'
+                test_df.to_sql(table_name2, uri, if_exists='append')
+
+                r = md.read_sql_table(table_name2, engine, chunk_size=4, index_col='id')
+                result = self.executor.execute_dataframe(r, concat=True)[0]
+                pd.testing.assert_frame_equal(result, test_df)
             finally:
                 engine.dispose()