feat(clickhouse): partition kwargs for compile and execution in to_pyarrow and to_pandas

Author: cpcloud

ibis/backends/clickhouse/__init__.py

@@ -308,7 +308,7 @@ def to_pyarrow_batches(
         params: Mapping[ir.Scalar, Any] | None = None,
         external_tables: Mapping[str, Any] | None = None,
         chunk_size: int = 1_000_000,
-        **_: Any,
+        **kwargs: Any,
     ) -> pa.ipc.RecordBatchReader:
         """Execute expression and return an iterator of pyarrow record batches.
 
@@ -328,6 +328,8 @@ def to_pyarrow_batches(
             External data
         chunk_size
             Maximum number of row to return in a single chunk
+        kwargs
+            Extra arguments passed directly to clickhouse-connect
 
         Returns
         -------
@@ -357,14 +359,17 @@ def to_pyarrow_batches(
         external_tables = self._collect_in_memory_tables(expr, external_tables)
         external_data = self._normalize_external_tables(external_tables)
 
-        def batcher(sql: str, *, schema: pa.Schema) -> Iterator[pa.RecordBatch]:
-            settings = {}
+        settings = kwargs.pop("settings", {})
 
-            # readonly != 1 means that the server setting is writable
-            if self.con.server_settings["max_block_size"].readonly != 1:
-                settings["max_block_size"] = chunk_size
+        # readonly != 1 means that the server setting is writable
+        if self.con.server_settings["max_block_size"].readonly != 1:
+            settings["max_block_size"] = chunk_size
+
+        def batcher(
+            sql: str, *, schema: pa.Schema, settings, **kwargs
+        ) -> Iterator[pa.RecordBatch]:
             with self.con.query_column_block_stream(
-                sql, external_data=external_data, settings=settings
+                sql, external_data=external_data, settings=settings, **kwargs
             ) as blocks:
                 yield from map(
                     partial(pa.RecordBatch.from_arrays, schema=schema), blocks
@@ -373,29 +378,34 @@ def batcher(sql: str, *, schema: pa.Schema) -> Iterator[pa.RecordBatch]:
         self._log(sql)
         schema = table.schema().to_pyarrow()
         return pa.ipc.RecordBatchReader.from_batches(
-            schema, batcher(sql, schema=schema)
+            schema, batcher(sql, schema=schema, settings=settings, **kwargs)
         )
 
     def execute(
         self,
         expr: ir.Expr,
         limit: str | None = "default",
+        params: Mapping[ir.Scalar, Any] | None = None,
         external_tables: Mapping[str, pd.DataFrame] | None = None,
         **kwargs: Any,
     ) -> Any:
         """Execute an expression."""
         import pandas as pd
 
         table = expr.as_table()
-        sql = self.compile(table, limit=limit, **kwargs)
+        sql = self.compile(table, params=params, limit=limit)
 
         schema = table.schema()
         self._log(sql)
 
         external_tables = self._collect_in_memory_tables(expr, external_tables)
         external_data = self._normalize_external_tables(external_tables)
         df = self.con.query_df(
-            sql, external_data=external_data, use_na_values=False, use_none=True
+            sql,
+            external_data=external_data,
+            use_na_values=False,
+            use_none=True,
+            **kwargs,
         )
 
         if df.empty:

ibis/backends/clickhouse/tests/test_client.py

@@ -425,3 +425,20 @@ def test_alias_column_ref(con):
     assert result.user_id.notnull().all()
     assert result.account_id.notnull().all()
     assert result.id_md5.notnull().all()
+
+
+@pytest.mark.parametrize("method_name", ["to_pandas", "to_pyarrow"])
+def test_query_cache(con, method_name):
+    t = con.table("functional_alltypes")
+    expr = t.count()
+
+    method = getattr(expr, method_name)
+
+    expected = method()
+    result = method(settings={"use_query_cache": True})
+
+    # test a bogus setting
+    with pytest.raises(ClickHouseDatabaseError):
+        method(settings={"ooze_query_cash": True})
+
+    assert result == expected