feat: add clickhouse data sink (#4850)

## Changes Made

<!-- Describe what changes were made and why. Include implementation
details if necessary. -->
Add the clickhouse data sink, and use it to writes the DataFrame to a
ClickHouse table
## Related Issues

<!-- Link to related GitHub issues, e.g., "Closes #123" -->

## Checklist

- [ ] Documented in API Docs (if applicable)
- [ ] Documented in User Guide (if applicable)
- [ ] If adding a new documentation page, doc is added to
`docs/mkdocs.yml` navigation
- [ ] Documentation builds and is formatted properly (tag @/ccmao1130
for docs review)

---------

Co-authored-by: Desmond Cheong <[email protected]>

Author: huleilei

daft/dataframe/dataframe.py

@@ -1594,6 +1594,57 @@ def write_turbopuffer(
         )
         return self.write_sink(sink)
 
+    @DataframePublicAPI
+    def write_clickhouse(
+        self,
+        table: str,
+        *,
+        host: str,
+        port: Optional[int] = None,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        database: Optional[str] = None,
+        client_kwargs: Optional[dict[str, Any]] = None,
+        write_kwargs: Optional[dict[str, Any]] = None,
+    ) -> "DataFrame":
+        """Writes the DataFrame to a ClickHouse table.
+
+        Args:
+            table: Name of the ClickHouse table to write to.
+            host: ClickHouse host.
+            port: ClickHouse port.
+            user: ClickHouse user.
+            password: ClickHouse password.
+            database: ClickHouse database.
+            client_kwargs: Optional dictionary of arguments to pass to the ClickHouse client constructor.
+            write_kwargs: Optional dictionary of arguments to pass to the ClickHouse write() method.
+
+        Examples:
+            >>> import daft
+            >>> df = daft.from_pydict({"a": [1, 2, 3, 4]})  # doctest: +SKIP
+            >>> df.write_clickhouse(table="", host="", port=8123, user="", password="")  # doctest: +SKIP
+            ╭────────────────────┬─────────────────────╮
+            │ total_written_rows ┆ total_written_bytes │
+            │ ---                ┆ ---                 │
+            │ Int64              ┆ Int64               │
+            ╞════════════════════╪═════════════════════╡
+            │ 4                  ┆ 32                  │
+            ╰────────────────────┴─────────────────────╯
+        """
+        from daft.io.clickhouse.clickhouse_data_sink import ClickHouseDataSink
+
+        sink = ClickHouseDataSink(
+            table,
+            host=host,
+            port=port,
+            user=user,
+            password=password,
+            database=database,
+            client_kwargs=client_kwargs,
+            write_kwargs=write_kwargs,
+        )
+        return self.write_sink(sink)
+
     ###
     # DataFrame operations
     ###

daft/io/clickhouse/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from daft.io.clickhouse.clickhouse_data_sink import ClickHouseDataSink
+
+__all__ = ["ClickHouseDataSink"]

daft/io/clickhouse/clickhouse_data_sink.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from clickhouse_connect import get_client
+from clickhouse_connect.driver.summary import QuerySummary
+
+from daft import Schema
+from daft.datatype import DataType
+from daft.io import DataSink
+from daft.io.sink import WriteResult
+from daft.recordbatch.micropartition import MicroPartition
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+class ClickHouseDataSink(DataSink[QuerySummary]):
+    def __init__(
+        self,
+        table: str,
+        *,
+        host: str,
+        port: int | None = None,
+        user: str | None = None,
+        password: str | None = None,
+        database: str | None = None,
+        client_kwargs: dict[str, Any] | None = None,
+        write_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        connection_params: dict[str, Any] = {
+            "host": host,
+        }
+
+        if port is not None:
+            connection_params["port"] = port
+        if user is not None:
+            connection_params["user"] = user
+        if password is not None:
+            connection_params["password"] = password
+        if database is not None:
+            connection_params["database"] = database
+
+        # Merge user-provided client_kwargs with explicit connection parameters
+        self._client_kwargs = {**(client_kwargs or {}), **connection_params}
+        self._table = table
+
+        self._write_kwargs = write_kwargs or {}
+
+        # Define the schema for the result of the write operation
+        self._result_schema = Schema._from_field_name_and_types(
+            [("total_written_rows", DataType.int64()), ("total_written_bytes", DataType.int64())]
+        )
+
+    def schema(self) -> Schema:
+        return self._result_schema
+
+    def write(self, micropartitions: Iterator[MicroPartition]) -> Iterator[WriteResult[QuerySummary]]:
+        """Writes to ClickHouse from the given micropartitions."""
+        # socket cannot be serialized, so we need to create a new client in write
+        ck_client = get_client(**self._client_kwargs)
+        try:
+            for micropartition in micropartitions:
+                df = micropartition.to_pandas()
+                bytes_written = df.memory_usage().sum()
+                rows_written = df.shape[0]
+
+                query_summary = ck_client.insert_df(self._table, df, **self._write_kwargs)
+                yield WriteResult(
+                    result=query_summary,
+                    bytes_written=bytes_written,
+                    rows_written=rows_written,
+                )
+        finally:
+            ck_client.close()
+
+    def finalize(self, write_results: list[WriteResult[QuerySummary]]) -> MicroPartition:
+        """Finish write to ClickHouse dataset. Returns a DataFrame with the stats of the dataset."""
+        from daft.dependencies import pa
+
+        total_written_rows = 0
+        total_written_bytes = 0
+
+        for write_result in write_results:
+            total_written_rows += write_result.rows_written
+            total_written_bytes += write_result.bytes_written
+
+        tbl = MicroPartition.from_pydict(
+            {
+                "total_written_rows": pa.array([total_written_rows], pa.int64()),
+                "total_written_bytes": pa.array([total_written_bytes], pa.int64()),
+            }
+        )
+        return tbl

pyproject.toml

@@ -22,9 +22,10 @@ readme = "README.rst"
 requires-python = ">=3.9"
 
 [project.optional-dependencies]
-all = ["daft[aws, azure, gcp, ray, pandas, numpy, iceberg, deltalake, spark, sql, unity]"]
+all = ["daft[aws, azure, gcp, ray, pandas, numpy, iceberg, deltalake, spark, sql, unity, clickhouse]"]
 aws = ["boto3"]
 azure = []
+clickhouse = ["clickhouse_connect"]
 deltalake = ["deltalake", "packaging"]
 gcp = []
 hudi = ["pyarrow >= 8.0.0"]

requirements-dev.txt

@@ -116,3 +116,6 @@ grpcio-status==1.67.0
 # ai
 vllm; platform_system == "Linux" and platform_machine == "x86_64" # for other systems, see install instructions: https://docs.vllm.ai/en/latest/getting_started/installation.html
 openai
+
+#clickhouse
+clickhouse-connect

tests/io/clickhouse/__init__.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import unittest
+from unittest.mock import MagicMock, Mock, patch
+
+import pandas as pd
+import pytest
+from clickhouse_connect.driver.summary import QuerySummary
+
+import daft
+from daft.dataframe.dataframe import DataFrame
+from daft.io.clickhouse.clickhouse_data_sink import ClickHouseDataSink
+
+
+class TestWriteClickHouse:
+    @pytest.fixture
+    def sample_data(self):
+        return daft.from_pydict(
+            {
+                "vector": [[1.1, 1.2], [0.2, 1.8]],
+                "lat": [45.5, 40.1],
+                "long": [-122.7, -74.1],
+            }
+        )
+
+    @patch.object(ClickHouseDataSink, "__init__", return_value=None)
+    @patch.object(DataFrame, "write_sink")
+    def test_write_clickhouse_minimal_params(self, mock_write_sink, mock_clickhouse_sink, sample_data):
+        """Test minimal parameters for write_clickhouse."""
+        sample_data.write_clickhouse(table="test_table", host="localhost")
+
+        mock_clickhouse_sink.assert_called_once_with(
+            "test_table",
+            host="localhost",
+            port=None,
+            user=None,
+            password=None,
+            database=None,
+            client_kwargs=None,
+            write_kwargs=None,
+        )
+        mock_write_sink.assert_called_once()
+
+    @patch.object(ClickHouseDataSink, "__init__", return_value=None)
+    @patch.object(DataFrame, "write_sink")
+    def test_write_clickhouse_all_params(self, mock_write_sink, mock_clickhouse_sink, sample_data):
+        """Test all parameters for write_clickhouse."""
+        sample_data.write_clickhouse(
+            table="test_table",
+            host="localhost",
+            port=8123,
+            user="user",
+            password="pass",
+            database="db",
+            client_kwargs={"timeout": 10},
+            write_kwargs={"batch_size": 1000},
+        )
+
+        mock_clickhouse_sink.assert_called_once_with(
+            "test_table",
+            host="localhost",
+            port=8123,
+            user="user",
+            password="pass",
+            database="db",
+            client_kwargs={"timeout": 10},
+            write_kwargs={"batch_size": 1000},
+        )
+        mock_write_sink.assert_called_once()
+
+    @patch.object(ClickHouseDataSink, "__init__", return_value=None)
+    @patch.object(DataFrame, "write_sink", side_effect=Exception("Write error"))
+    def test_write_clickhouse_write_error(self, mock_write_sink, mock_clickhouse_sink, sample_data):
+        """Test write error for write_clickhouse."""
+        with pytest.raises(Exception, match="Write error"):
+            sample_data.write_clickhouse(table="test_table", host="localhost")
+
+
+class TestClickHouseDataSink:
+    """Tests for ClickHouseDataSink class."""
+
+    @pytest.fixture
+    def clickhouse_sink(self):
+        return ClickHouseDataSink(
+            table="test_table", host="localhost", port=8123, user="user", password="pass", database="db"
+        )
+
+    def test_initialization(self, clickhouse_sink):
+        assert clickhouse_sink._table == "test_table"
+        assert clickhouse_sink._client_kwargs["host"] == "localhost"
+        assert clickhouse_sink._client_kwargs["port"] == 8123
+        assert clickhouse_sink._client_kwargs["user"] == "user"
+        assert clickhouse_sink._client_kwargs["password"] == "pass"
+        assert clickhouse_sink._client_kwargs["database"] == "db"
+
+    @pytest.fixture
+    def mock_client(self):
+        with patch("daft.io.clickhouse.clickhouse_data_sink.get_client") as mock:  # 修正mock路径
+            client = MagicMock()
+            client.insert_df.return_value = QuerySummary(
+                summary={"written_rows": 3, "written_bytes": 128, "query_id": "test_query"}
+            )
+            client.command.return_value = ("22.8.0", "UTC")
+            mock.return_value = client
+            yield client
+
+    def test_clickhouse_sink_write(self, mock_client):
+        sink = ClickHouseDataSink(
+            table="test_table",
+            host="localhost",
+            client_kwargs={"settings": {"async_insert": 1}},
+            write_kwargs={"column_names": ["col1"]},
+        )
+
+        mp = daft.recordbatch.MicroPartition.from_pydict({"col1": [1, 2, 3]})
+        results = list(sink.write(iter([mp])))
+
+        mock_client.insert_df.assert_called_once_with("test_table", unittest.mock.ANY, column_names=["col1"])
+        pd.testing.assert_frame_equal(mock_client.insert_df.call_args[0][1], mp.to_pandas())
+
+        assert len(results) == 1
+        assert isinstance(results[0].result, QuerySummary)
+
+    def test_finalize_statistics(self):
+        sink = ClickHouseDataSink(table="test", host="localhost")
+        mock_results = [Mock(rows_written=10, bytes_written=100), Mock(rows_written=20, bytes_written=200)]
+
+        result = sink.finalize(mock_results)
+        assert result.to_pydict() == {"total_written_rows": [30], "total_written_bytes": [300]}
+
+    def test_client_cleanup(self, mock_client):
+        sink = ClickHouseDataSink(table="test", host="localhost")
+        mp = daft.recordbatch.MicroPartition.from_pydict({"col1": [1]})
+
+        mock_client.insert_df.side_effect = Exception("DB error")
+        with pytest.raises(Exception):
+            list(sink.write(iter([mp])))
+
+        mock_client.close.assert_called_once()