SQLAlchemy: Add documentation and tests for usage with Dask [WIP]

Author: amotl

CHANGES.txt

@@ -7,6 +7,7 @@ Unreleased
 
 - SQLAlchemy: Added ``insert_bulk`` fast-path ``INSERT`` method for pandas, in
   order to support efficient batch inserts using CrateDB's bulk operations endpoint.
+- SQLAlchemy: Add documentation and software tests for usage with Dask [WIP]
 
 
 2023/04/18 0.31.1

docs/by-example/sqlalchemy/dataframe.rst

@@ -5,6 +5,12 @@
 SQLAlchemy: DataFrame operations
 ================================
 
+.. rubric:: Table of Contents
+
+.. contents::
+   :local:
+
+
 About
 =====
 
@@ -29,10 +35,13 @@ The :ref:`pandas I/O subsystem <pandas:api.io>` for `relational databases`_
 using `SQL`_ is based on `SQLAlchemy`_.
 
 
-.. rubric:: Table of Contents
+Compatibility notes
+===================
 
-.. contents::
-   :local:
+.. NOTE::
+
+    Please note that DataFrame support for pandas and Dask is only validated
+    with Python 3.8 and higher, and SQLAlchemy 1.4 and higher.
 
 
 Efficient ``INSERT`` operations with pandas

setup.py

@@ -68,6 +68,7 @@ def read(path):
               'zope.testrunner>=5,<7',
               'zc.customdoctests>=1.0.1,<2',
               'createcoverage>=1,<2',
+              'dask',
               'stopit>=1.1.2,<2',
               'flake8>=4,<7',
               'pandas',

src/crate/client/sqlalchemy/tests/bulk_test.py

@@ -207,3 +207,46 @@ def test_bulk_save_pandas(self, mock_cursor):
 
         # Verify number of batches.
         self.assertEqual(effective_op_count, OPCOUNT)
+
+    @skipIf(sys.version_info < (3, 8), "SQLAlchemy/Dask is not supported on Python <3.8")
+    @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 is not supported by pandas")
+    @patch('crate.client.connection.Cursor', mock_cursor=FakeCursor)
+    def test_bulk_save_dask(self, mock_cursor):
+        """
+        Verify bulk INSERT with Dask.
+        """
+        import dask.dataframe as dd
+        from pandas._testing import makeTimeDataFrame
+        from crate.client.sqlalchemy.support import insert_bulk
+
+        # 42 records / 4 partitions means each partition has a size of 10.5 elements.
+        # Because the chunk size 8 is slightly smaller than 10, the partition will not
+        # fit into it, so two batches will be emitted to the database for each data
+        # partition. 4 partitions * 2 batches = 8 insert operations will be emitted.
+        INSERT_RECORDS = 42
+        NPARTITIONS = 4
+        CHUNK_SIZE = 8
+        OPCOUNT = math.ceil(INSERT_RECORDS / NPARTITIONS / CHUNK_SIZE) * NPARTITIONS
+
+        # Create a DataFrame to feed into the database.
+        df = makeTimeDataFrame(nper=INSERT_RECORDS, freq="S")
+        ddf = dd.from_pandas(df, npartitions=NPARTITIONS)
+
+        dburi = "crate://localhost:4200"
+        retval = ddf.to_sql(
+            name="test-testdrive",
+            uri=dburi,
+            if_exists="replace",
+            index=False,
+            chunksize=CHUNK_SIZE,
+            method=insert_bulk,
+        )
+        self.assertIsNone(retval)
+
+        # Each of the insert operation incurs another call to the cursor object. This is probably
+        # the initial connection from the DB-API driver, to inquire the database version.
+        # This compensation formula has been determined empirically / by educated guessing.
+        effective_op_count = (mock_cursor.call_count - 2 * NPARTITIONS) - 2
+
+        # Verify number of batches.
+        self.assertEqual(effective_op_count, OPCOUNT)