Repository: dask/dask

Issue:
read_sql() and to_sql()?
It would be wonderful to have both read_sql() and to_sql(), analogous to pandas for those times when you have to clean up inconveniently large datasets. I saw [http://stackoverflow.com/questions/33129839/choosing-a-framework-for-larger-than-memory-data-analysis-with-python](url) but it appears that no issue had been opened yet



Discussion:
Yup.  I agree that this would be great to have..  For read_sql we would have to include a mandatory order by column to the API in order to cleanly pull out partitions.  See http://www.dofactory.com/sql/order-by-offset-fetch 

This will only work efficiently if there is some indexed column and the user knows what it is.

Something more to keep in mind (not that I've tested any of these flavours - I'm stuck with MySQL): [http://blog.jooq.org/2014/06/09/stop-trying-to-emulate-sql-offset-pagination-with-your-in-house-db-framework/](url)

As to indexed column and efficiency ... I'd argue it'd be a great start to expect the user to accommodate this need or suffer the consequences, i.e. low performance

@mrocklin on 2nd thought: why would that be different from read_csv, or is your comment about the efficiency of a read_sql process?

We need to ensure that several separate queries properly partition the data into disjoint subsets that cover the entire table.  I think that the only way to do this is with order-by/limit-offset but would be happy to learn of alternatives.  Additionally, sorting the dataset on an index would align well with other dask.dataframe operations

It now occurs to me that we don't really want limit/offset, we want to find percentiles on a column and then filter with a where clause for each of the partitions.  This is more in line with the partition by index-value approach.

This should be fairly easy to implement with `pd.read_sql`, `dask.delayed` and `dask.dataframe.from_delayed`.

Pseudo code would be something like the following:

``` python
def read_sql_table(table_name, con, schema=None, index_col=None, coerce_float=True, parse_dates=None, columns=None, npartitions=None):
    if npartitions is None:
        length = pd.read_sql('select count(*) from table')
        npartitions = length // some_reasonable_chunksize
    divisions = pd.read_sql_query('select index_col and something with ntile')
    parts = [pd.read_sql_query('select columns from table where index_col >= divisions[i] and index_col < divisions[i + 1]') for i in range(len(divisions) - 1)]
    head = pd.read_sql('select columns from table limit 5')
    return dd.from_delayed(parts, divisions, head)
```

We can probably test against sqlite.  We should probably have pretty loud text warning that the index_column should really be indexed.


Resolving patch:
diff --git a/dask/dataframe/__init__.py b/dask/dataframe/__init__.py
--- a/dask/dataframe/__init__.py
+++ b/dask/dataframe/__init__.py
@@ -3,7 +3,7 @@
 from .core import (DataFrame, Series, Index, _Frame, map_partitions,
                    repartition, to_delayed)
 from .io import (from_array, from_pandas, from_bcolz,
-                 from_dask_array, read_hdf,
+                 from_dask_array, read_hdf, read_sql_table,
                  from_delayed, read_csv, to_csv, read_table,
                  demo, to_hdf, to_records, to_bag)
 from .optimize import optimize
diff --git a/dask/dataframe/io/__init__.py b/dask/dataframe/io/__init__.py
--- a/dask/dataframe/io/__init__.py
+++ b/dask/dataframe/io/__init__.py
@@ -5,6 +5,7 @@
                  from_delayed, dataframe_from_ctable, to_bag, to_records)
 from .csv import read_csv, to_csv, read_table
 from .hdf import read_hdf, to_hdf
+from .sql import read_sql_table
 from . import demo
 try:
     from .parquet import read_parquet, to_parquet
diff --git a/dask/dataframe/io/sql.py b/dask/dataframe/io/sql.py
new file mode 100644
--- /dev/null
+++ b/dask/dataframe/io/sql.py
@@ -0,0 +1,131 @@
+import numpy as np
+import pandas as pd
+import six
+
+from ... import delayed
+from .io import from_delayed
+
+
+def read_sql_table(table, uri, index_col, divisions=None, npartitions=None,
+                   limits=None, columns=None, bytes_per_chunk=256 * 2**20,
+                   **kwargs):
+    """
+    Create dataframe from an SQL table.
+
+    If neither divisions or npartitions is given, the memory footprint of the
+    first five rows will be determined, and partitions of size ~256MB will
+    be used.
+
+    Parameters
+    ----------
+    table : string or sqlalchemy expression
+        Select columns from here.
+    uri : string
+        Full sqlalchemy URI for the database connection
+    index_col : string
+        Column which becomes the index, and defines the partitioning. Should
+        be a indexed column in the SQL server, and numerical.
+        Could be a function to return a value, e.g.,
+        ``sql.func.abs(sql.column('value')).label('abs(value)')``.
+        Labeling columns created by functions or arithmetic operations is
+        required.
+    divisions: sequence
+        Values of the index column to split the table by.
+    npartitions : int
+        Number of partitions, if divisions is not given. Will split the values
+        of the index column linearly between limits, if given, or the column
+        max/min.
+    limits: 2-tuple or None
+        Manually give upper and lower range of values for use with npartitions;
+        if None, first fetches max/min from the DB. Upper limit, if
+        given, is inclusive.
+    columns : list of strings or None
+        Which columns to select; if None, gets all; can include sqlalchemy
+        functions, e.g.,
+        ``sql.func.abs(sql.column('value')).label('abs(value)')``.
+        Labeling columns created by functions or arithmetic operations is
+        recommended.
+    bytes_per_chunk: int
+        If both divisions and npartitions is None, this is the target size of
+        each partition, in bytes
+    kwargs : dict
+        Additional parameters to pass to `pd.read_sql()`
+
+    Returns
+    -------
+    dask.dataframe
+
+    Examples
+    --------
+    >>> df = dd.read_sql('accounts', 'sqlite:///path/to/bank.db',
+    ...                  npartitions=10, index_col='id')  # doctest: +SKIP
+    """
+    import sqlalchemy as sa
+    from sqlalchemy import sql
+    from sqlalchemy.sql import elements
+    if index_col is None:
+        raise ValueError("Must specify index column to partition on")
+    engine = sa.create_engine(uri)
+    meta = sa.MetaData()
+    if isinstance(table, six.string_types):
+        table = sa.Table(table, meta, autoload=True, autoload_with=engine)
+
+    index = (table.columns[index_col] if isinstance(index_col, six.string_types)
+             else index_col)
+    if not isinstance(index_col, six.string_types + (elements.Label,)):
+        raise ValueError('Use label when passing an SQLAlchemy instance'
+                         ' as the index (%s)' % index)
+    if divisions and npartitions:
+        raise TypeError('Must supply either divisions or npartitions, not both')
+
+    columns = ([(table.columns[c] if isinstance(c, six.string_types) else c)
+                for c in columns]
+               if columns else list(table.columns))
+    if index_col not in columns:
+        columns.append(table.columns[index_col]
+                       if isinstance(index_col, six.string_types)
+                       else index_col)
+
+    if isinstance(index_col, six.string_types):
+        kwargs['index_col'] = index_col
+    else:
+        # function names get pandas auto-named
+        kwargs['index_col'] = index_col.name
+
+    q = sql.select(columns).limit(5).select_from(table)
+    head = pd.read_sql(q, engine, **kwargs)
+
+    if divisions is None:
+        if limits is None:
+            # calculate max and min for given index
+            q = sql.select([sql.func.max(index), sql.func.min(index)]
+                           ).select_from(table)
+            minmax = pd.read_sql(q, engine)
+            maxi, mini = minmax.iloc[0]
+            dtype = minmax.dtypes['max_1']
+        else:
+            mini, maxi = limits
+            dtype = pd.Series(limits).dtype
+        if npartitions is None:
+            q = sql.select([sql.func.count(index)]).select_from(table)
+            count = pd.read_sql(q, engine)['count_1'][0]
+            bytes_per_row = (head.memory_usage(deep=True, index=True)).sum() / 5
+            npartitions = round(count * bytes_per_row / bytes_per_chunk)
+        if dtype.kind == "M":
+            divisions = pd.date_range(
+                start=mini, end=maxi, freq='%iS' % (
+                    (maxi - mini) / npartitions).total_seconds()).tolist()
+            divisions[0] = mini
+            divisions[-1] = maxi
+        else:
+            divisions = np.linspace(mini, maxi, npartitions + 1).tolist()
+
+    parts = []
+    lowers, uppers = divisions[:-1], divisions[1:]
+    for i, (lower, upper) in enumerate(zip(lowers, uppers)):
+        cond = index <= upper if i == len(lowers) - 1 else index < upper
+        q = sql.select(columns).where(sql.and_(index >= lower, cond)
+                                      ).select_from(table)
+        parts.append(delayed(pd.read_sql)(q, uri, **kwargs))
+
+    return from_delayed(parts, head, divisions=divisions)


Test patch:
diff --git a/dask/dataframe/io/tests/test_sql.py b/dask/dataframe/io/tests/test_sql.py
new file mode 100644
--- /dev/null
+++ b/dask/dataframe/io/tests/test_sql.py
@@ -0,0 +1,143 @@
+from __future__ import (print_function, division, absolute_import,
+                        unicode_literals)
+
+import io
+import pytest
+
+from dask.dataframe.io.sql import read_sql_table
+from dask.utils import tmpfile
+from dask.dataframe.utils import assert_eq
+pd = pytest.importorskip('pandas')
+dd = pytest.importorskip('dask.dataframe')
+pytest.importorskip('sqlalchemy')
+pytest.importorskip('sqlite3')
+
+
+data = """
+name,number,age,negish
+Alice,0,33,-5
+Bob,1,40,-3
+Chris,2,22,3
+Dora,3,16,5
+Edith,4,53,0
+Francis,5,30,0
+Garreth,6,20,0
+"""
+
+df = pd.read_csv(io.StringIO(data), index_col='number')
+
+
+@pytest.yield_fixture
+def db():
+    with tmpfile() as f:
+        uri = 'sqlite:///%s' % f
+        df.to_sql('test', uri, index=True, if_exists='replace')
+        yield uri
+
+
+def test_simple(db):
+    # single chunk
+    data = read_sql_table('test', db, npartitions=2, index_col='number'
+                          ).compute()
+    assert (data.name == df.name).all()
+    assert data.index.name == 'number'
+    assert_eq(data, df)
+
+
+def test_npartitions(db):
+    data = read_sql_table('test', db, columns=list(df.columns), npartitions=2,
+                          index_col='number')
+    assert len(data.divisions) == 3
+    assert (data.name.compute() == df.name).all()
+    data = read_sql_table('test', db, columns=['name'], npartitions=6,
+                          index_col="number")
+    assert_eq(data, df[['name']])
+
+
+def test_divisions(db):
+    data = read_sql_table('test', db, columns=['name'], divisions=[0, 2, 4],
+                          index_col="number")
+    assert data.divisions == (0, 2, 4)
+    assert data.index.max().compute() == 4
+    assert_eq(data, df[['name']][df.index <= 4])
+
+
+def test_division_or_partition(db):
+    with pytest.raises(TypeError):
+        read_sql_table('test', db, columns=['name'], index_col="number",
+                       divisions=[0, 2, 4], npartitions=3)
+
+    out = read_sql_table('test', db, index_col="number", bytes_per_chunk=100)
+    m = out.map_partitions(lambda d: d.memory_usage(
+        deep=True, index=True).sum()).compute()
+    assert (50 < m).all() and (m < 200).all()
+    assert_eq(out, df)
+
+
+def test_range(db):
+    data = read_sql_table('test', db, npartitions=2, index_col='number',
+                          limits=[1, 4])
+    assert data.index.min().compute() == 1
+    assert data.index.max().compute() == 4
+
+
+def test_datetimes():
+    import datetime
+    now = datetime.datetime.now()
+    d = datetime.timedelta(seconds=1)
+    df = pd.DataFrame({'a': list('ghjkl'), 'b': [now + i * d
+                                                 for i in range(2, -3, -1)]})
+    with tmpfile() as f:
+        uri = 'sqlite:///%s' % f
+        df.to_sql('test', uri, index=False, if_exists='replace')
+        data = read_sql_table('test', uri, npartitions=2, index_col='b')
+        assert data.index.dtype.kind == "M"
+        assert data.divisions[0] == df.b.min()
+        df2 = df.set_index('b')
+        assert_eq(data.map_partitions(lambda x: x.sort_index()),
+                  df2.sort_index())
+
+
+def test_with_func(db):
+    from sqlalchemy import sql
+    index = sql.func.abs(sql.column('negish')).label('abs')
+
+    # function for the index, get all columns
+    data = read_sql_table('test', db, npartitions=2, index_col=index)
+    assert data.divisions[0] == 0
+    part = data.get_partition(0).compute()
+    assert (part.index == 0).all()
+
+    # now an arith op for one column too; it's name will be 'age'
+    data = read_sql_table('test', db, npartitions=2, index_col=index,
+                          columns=[index, -sql.column('age')])
+    assert (data.age.compute() < 0).all()
+
+    # a column that would have no name, give it a label
+    index = (-sql.column('negish')).label('index')
+    data = read_sql_table('test', db, npartitions=2, index_col=index,
+                          columns=['negish', 'age'])
+    d = data.compute()
+    assert (-d.index == d['negish']).all()
+
+
+def test_no_nameless_index(db):
+    from sqlalchemy import sql
+    index = (-sql.column('negish'))
+    with pytest.raises(ValueError):
+        read_sql_table('test', db, npartitions=2, index_col=index,
+                       columns=['negish', 'age', index])
+
+    index = sql.func.abs(sql.column('negish'))
+
+    # function for the index, get all columns
+    with pytest.raises(ValueError):
+        read_sql_table('test', db, npartitions=2, index_col=index)
+
+
+def test_select_from_select(db):
+    from sqlalchemy import sql
+    s1 = sql.select([sql.column('number'), sql.column('name')]
+                    ).select_from(sql.table('test'))
+    out = read_sql_table(s1, db, npartitions=2, index_col='number')
+    assert_eq(out, df[['name']])