Repository: modin-project/modin

Issue:
Can't pickle _thread.RLock objects when using read_sql
### System information
- **OS Platform and Distribution**: Ubuntu 18.04 (docker python-stretch container)
- **Modin installed from**: binary (pip)
- **Modin version**: 0.5.3
- **SQLAlchemy version**: 1.3.5
- **Psycopg2 version**: 2.8.3
- **Python version**: 3.7.3
- **DB**: Amazon RedShift

### The Problem
Function `read_sql` fails because RLock objects can't be pickled

### Source Code
```python
import sqlalchemy as sa
import modin.pandas as pd

config = {'foo': 'bar'}
conn_str = sa.engine.url.URL(**config, drivername='postgresql+psycopg2')
engine = sa.create_engine(conn_str)

query = open('some_query.sql').read()
with engine.connect() as conn, conn.begin():
    data = pd.read_sql(query, conn, parse_dates=['install_date'])
```

### Logs
```
2019-06-26 14:27:38,787	WARNING worker.py:1341 -- WARNING: Not updating worker name since `setproctitle` is not installed. Install this with `pip install setproctitle` (or ray[debug]) to enable monitoring of worker processes.
2019-06-26 14:27:38,788	INFO node.py:497 -- Process STDOUT and STDERR is being redirected to /tmp/ray/session_2019-06-26_14-27-38_788074_1/logs.
2019-06-26 14:27:38,892	INFO services.py:409 -- Waiting for redis server at 127.0.0.1:53765 to respond...
2019-06-26 14:27:39,038	INFO services.py:409 -- Waiting for redis server at 127.0.0.1:37958 to respond...
2019-06-26 14:27:39,041	INFO services.py:806 -- Starting Redis shard with 10.0 GB max memory.
2019-06-26 14:27:39,077	INFO node.py:511 -- Process STDOUT and STDERR is being redirected to /tmp/ray/session_2019-06-26_14-27-38_788074_1/logs.
2019-06-26 14:27:39,086	INFO services.py:1441 -- Starting the Plasma object store with 80.0 GB memory using /dev/shm.
2019-06-26 14:27:39,098	WARNING services.py:909 -- Failed to start the reporter. The reporter requires 'pip install psutil'.
2019-06-26 15:03:14,196	WARNING worker.py:342 -- WARNING: Falling back to serializing objects of type <class '_thread.RLock'> by using pickle. This may be inefficient.
2019-06-26 15:03:14,211	WARNING worker.py:402 -- WARNING: Serializing the class <class 'sqlalchemy.engine.base.Connection'> failed, so are are falling back to cloudpickle.
WARNING: Not monitoring node memory since `psutil` is not installed. Install this with `pip install psutil` (or ray[debug]) to enable debugging of memory-related crashes.
```

### Traceback
```
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/site-packages/ray/worker.py", line 383, in put_object
    self.store_and_register(object_id, value)
  File "/usr/local/lib/python3.7/site-packages/ray/worker.py", line 317, in store_and_register
    self.task_driver_id))
  File "/usr/local/lib/python3.7/site-packages/ray/utils.py", line 475, in _wrapper
    return orig_attr(*args, **kwargs)
  File "pyarrow/_plasma.pyx", line 496, in pyarrow._plasma.PlasmaClient.put
  File "pyarrow/serialization.pxi", line 355, in pyarrow.lib.serialize
  File "pyarrow/serialization.pxi", line 150, in pyarrow.lib.SerializationContext._serialize_callback
  File "/usr/local/lib/python3.7/site-packages/ray/cloudpickle/cloudpickle.py", line 952, in dumps
    cp.dump(obj)
  File "/usr/local/lib/python3.7/site-packages/ray/cloudpickle/cloudpickle.py", line 267, in dump
    return Pickler.dump(self, obj)
  File "/usr/local/lib/python3.7/pickle.py", line 437, in dump
    self.save(obj)
  File "/usr/local/lib/python3.7/pickle.py", line 524, in save
    rv = reduce(self.proto)
TypeError: can't pickle _thread.RLock objects

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "src/main.py", line 55, in <module>
    start_date=start_date)
  File "/src/preparation/data_source.py", line 16, in import_data
    data = self._load_data(app_ids, granularity, f'{date.today()}_{self._source}_{granularity}.csv')
  File "/src/preparation/data_source.py", line 27, in _load_data
    data = pd.read_sql(query, conn, parse_dates=['install_date'])
  File "/usr/local/lib/python3.7/site-packages/modin/pandas/io.py", line 324, in read_sql
    return DataFrame(query_compiler=BaseFactory.read_sql(**kwargs))
  File "/usr/local/lib/python3.7/site-packages/modin/data_management/factories.py", line 155, in read_sql
    return cls._determine_engine()._read_sql(**kwargs)
  File "/usr/local/lib/python3.7/site-packages/modin/data_management/factories.py", line 159, in _read_sql
    return cls.io_cls.read_sql(**kwargs)
  File "/usr/local/lib/python3.7/site-packages/modin/engines/ray/generic/io.py", line 805, in read_sql
    num_return_vals=num_parts + 1,
  File "/usr/local/lib/python3.7/site-packages/ray/remote_function.py", line 152, in _remote
    return invocation(args, kwargs)
  File "/usr/local/lib/python3.7/site-packages/ray/remote_function.py", line 143, in invocation
    resources=resources)
  File "/usr/local/lib/python3.7/site-packages/ray/worker.py", line 636, in submit_task
    args_for_raylet.append(put(arg))
  File "/usr/local/lib/python3.7/site-packages/ray/worker.py", line 2227, in put
    worker.put_object(object_id, value)
  File "/usr/local/lib/python3.7/site-packages/ray/worker.py", line 403, in put_object
    self.store_and_register(object_id, value)
  File "/usr/local/lib/python3.7/site-packages/ray/worker.py", line 317, in store_and_register
    self.task_driver_id))
  File "/usr/local/lib/python3.7/site-packages/ray/utils.py", line 475, in _wrapper
    return orig_attr(*args, **kwargs)
  File "pyarrow/_plasma.pyx", line 496, in pyarrow._plasma.PlasmaClient.put
  File "pyarrow/serialization.pxi", line 355, in pyarrow.lib.serialize
  File "pyarrow/serialization.pxi", line 150, in pyarrow.lib.SerializationContext._serialize_callback
  File "/usr/local/lib/python3.7/site-packages/ray/cloudpickle/cloudpickle.py", line 952, in dumps
    cp.dump(obj)
  File "/usr/local/lib/python3.7/site-packages/ray/cloudpickle/cloudpickle.py", line 267, in dump
    return Pickler.dump(self, obj)
  File "/usr/local/lib/python3.7/pickle.py", line 437, in dump
    self.save(obj)
  File "/usr/local/lib/python3.7/pickle.py", line 549, in save
    self.save_reduce(obj=obj, *rv)
  File "/usr/local/lib/python3.7/pickle.py", line 662, in save_reduce
    save(state)
  File "/usr/local/lib/python3.7/pickle.py", line 504, in save
    f(self, obj) # Call unbound method with explicit self
  File "/usr/local/lib/python3.7/pickle.py", line 856, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/local/lib/python3.7/pickle.py", line 882, in _batch_setitems
    save(v)
  File "/usr/local/lib/python3.7/pickle.py", line 549, in save
    self.save_reduce(obj=obj, *rv)
  File "/usr/local/lib/python3.7/pickle.py", line 662, in save_reduce
    save(state)
  File "/usr/local/lib/python3.7/pickle.py", line 504, in save
    f(self, obj) # Call unbound method with explicit self
  File "/usr/local/lib/python3.7/pickle.py", line 856, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/local/lib/python3.7/pickle.py", line 882, in _batch_setitems
    save(v)
  File "/usr/local/lib/python3.7/pickle.py", line 549, in save
    self.save_reduce(obj=obj, *rv)
  File "/usr/local/lib/python3.7/pickle.py", line 662, in save_reduce
    save(state)
  File "/usr/local/lib/python3.7/pickle.py", line 504, in save
    f(self, obj) # Call unbound method with explicit self
  File "/usr/local/lib/python3.7/pickle.py", line 856, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/local/lib/python3.7/pickle.py", line 882, in _batch_setitems
    save(v)
  File "/usr/local/lib/python3.7/pickle.py", line 524, in save
    rv = reduce(self.proto)
TypeError: can't pickle _thread._local objects
```


Discussion:
Hi @lisovskey thanks for posting! I will look into this. We should not be trying to pickle anything from `_thread` in `read_sql`, so it looks like we are making some assumptions about the nature of the connection object. This is definitely a bug, thanks for posting!

Resolving patch:
diff --git a/modin/engines/ray/generic/io.py b/modin/engines/ray/generic/io.py
--- a/modin/engines/ray/generic/io.py
+++ b/modin/engines/ray/generic/io.py
@@ -808,10 +808,17 @@ def read_sql(cls, sql, con, index_col=None, **kwargs):
         if cls.read_sql_remote_task is None:
             return super(RayIO, cls).read_sql(sql, con, index_col=index_col, **kwargs)
 
-        row_cnt_query = "SELECT COUNT(*) FROM ({} as foo)".format(sql)
+        import sqlalchemy as sa
+
+        # In the case that we are given a SQLAlchemy Connection or Engine, the objects
+        # are not pickleable. We have to convert it to the URL string and connect from
+        # each of the workers.
+        if isinstance(con, (sa.engine.Engine, sa.engine.Connection)):
+            con = repr(con.engine.url)
+        row_cnt_query = "SELECT COUNT(*) FROM ({}) as foo".format(sql)
         row_cnt = pandas.read_sql(row_cnt_query, con).squeeze()
         cols_names_df = pandas.read_sql(
-            "SELECT * FROM ({} as foo) LIMIT 0".format(sql), con, index_col=index_col
+            "SELECT * FROM ({}) as foo LIMIT 0".format(sql), con, index_col=index_col
         )
         cols_names = cols_names_df.columns
         num_parts = cls.frame_mgr_cls._compute_num_partitions()
@@ -820,7 +827,7 @@ def read_sql(cls, sql, con, index_col=None, **kwargs):
         limit = math.ceil(row_cnt / num_parts)
         for part in range(num_parts):
             offset = part * limit
-            query = "SELECT * FROM ({} as foo) LIMIT {} OFFSET {}".format(
+            query = "SELECT * FROM ({}) as foo LIMIT {} OFFSET {}".format(
                 sql, limit, offset
             )
             partition_id = cls.read_sql_remote_task._remote(


Test patch:
diff --git a/modin/pandas/test/test_io.py b/modin/pandas/test/test_io.py
--- a/modin/pandas/test/test_io.py
+++ b/modin/pandas/test/test_io.py
@@ -11,6 +11,7 @@
 import pyarrow as pa
 import os
 import shutil
+import sqlalchemy as sa
 import sys
 
 from .utils import df_equals
@@ -327,7 +328,6 @@ def _sql_connection(filename, table=""):
         if os.path.exists(filename):
             os.remove(filename)
         filenames.append(filename)
-
         # Create connection and, if needed, table
         conn = "sqlite:///{}".format(filename)
         if table:
@@ -579,6 +579,20 @@ def test_from_sql(make_sql_connection):
     with pytest.warns(UserWarning):
         pd.read_sql_table(table, conn)
 
+    # Test SQLAlchemy engine
+    conn = sa.create_engine(conn)
+    pandas_df = pandas.read_sql(query, conn)
+    modin_df = pd.read_sql(query, conn)
+
+    assert modin_df_equals_pandas(modin_df, pandas_df)
+
+    # Test SQLAlchemy Connection
+    conn = conn.connect()
+    pandas_df = pandas.read_sql(query, conn)
+    modin_df = pd.read_sql(query, conn)
+
+    assert modin_df_equals_pandas(modin_df, pandas_df)
+
 
 @pytest.mark.skip(reason="No SAS write methods in Pandas")
 def test_from_sas():