From 6cff2d66b0337320025759f8d8b564a90a261ba5 Mon Sep 17 00:00:00 2001 From: GabrielBBaldez <130607246+GabrielBBaldez@users.noreply.github.com> Date: Mon, 29 Jun 2026 19:27:23 -0300 Subject: [PATCH] Add a ClickHouse connector with batched writes --- .../3.connector/11.clickhouse.md | 67 +++++ .../3.connector/index.rst | 1 + .../3.connector/11.clickhouse.md | 67 +++++ .../3.connector/index.rst | 3 +- .../geaflow-dsl-connector-clickhouse/pom.xml | 75 +++++ .../clickhouse/ClickHouseConfigKeys.java | 81 +++++ .../clickhouse/ClickHouseTableConnector.java | 46 +++ .../clickhouse/ClickHouseTableSink.java | 144 +++++++++ .../clickhouse/ClickHouseTableSource.java | 280 ++++++++++++++++++ .../clickhouse/util/ClickHouseUtils.java | 111 +++++++ ...e.geaflow.dsl.connector.api.TableConnector | 20 ++ .../ClickHouseTableConnectorTest.java | 280 ++++++++++++++++++ .../geaflow-dsl/geaflow-dsl-connector/pom.xml | 1 + 13 files changed, 1175 insertions(+), 1 deletion(-) create mode 100644 docs/docs-cn/source/5.application-development/3.connector/11.clickhouse.md create mode 100644 docs/docs-en/source/5.application-development/3.connector/11.clickhouse.md create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/pom.xml create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseConfigKeys.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnector.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSink.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSource.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/util/ClickHouseUtils.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/resources/META-INF/services/org.apache.geaflow.dsl.connector.api.TableConnector create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/test/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnectorTest.java diff --git a/docs/docs-cn/source/5.application-development/3.connector/11.clickhouse.md b/docs/docs-cn/source/5.application-development/3.connector/11.clickhouse.md new file mode 100644 index 000000000..21ec255ba --- /dev/null +++ b/docs/docs-cn/source/5.application-development/3.connector/11.clickhouse.md @@ -0,0 +1,67 @@ +# ClickHouse Connector介绍 +ClickHouse Connector支持读和写。 + +写入时会先将数据缓冲，再按批量（batch）写入ClickHouse。由于ClickHouse是列式存储、针对批量写入做了优化，批量写入相比逐行写入（通过通用JDBC Connector）有显著更高的吞吐。批次大小可配置。 + +读取时可以按某个数值列将读取拆分为多个分区，从而并行读取。 + +## 语法 + +```sql +CREATE TABLE clickhouse_table ( + id BIGINT, + name VARCHAR, + score DOUBLE +) WITH ( + type='clickhouse', + geaflow.dsl.clickhouse.url = 'jdbc:clickhouse://localhost:8123/default', + geaflow.dsl.clickhouse.username = 'default', + geaflow.dsl.clickhouse.password = '', + geaflow.dsl.clickhouse.table.name = 'source_table' +); +``` + +## 参数 + +| 参数名 | 是否必须 | 描述 | +| -------- |------|---------------------------------------------------| +| geaflow.dsl.clickhouse.url | 是 | The ClickHouse JDBC url, e.g. jdbc:clickhouse://host:8123/database. | +| geaflow.dsl.clickhouse.table.name | 是 | The ClickHouse table name. | +| geaflow.dsl.clickhouse.username | 否 | The ClickHouse username, default 'default'. | +| geaflow.dsl.clickhouse.password | 否 | The ClickHouse password, default empty. | +| geaflow.dsl.clickhouse.write.batch.size | 否 | The number of rows buffered before a batch is flushed to ClickHouse. Larger batches give higher write throughput. Default 1000. | +| geaflow.dsl.clickhouse.driver | 否 | The ClickHouse JDBC driver class, default com.clickhouse.jdbc.ClickHouseDriver. | +| geaflow.dsl.clickhouse.partition.num | 否 | The number of source read partitions, default 1. | +| geaflow.dsl.clickhouse.partition.column | 否 | The column used to split the source read into partitions. Default value is 'id'. | +| geaflow.dsl.clickhouse.partition.lowerbound | 否 | The lowerbound used to decide the partition stride, not for filtering the rows in the table. | +| geaflow.dsl.clickhouse.partition.upperbound | 否 | The upperbound used to decide the partition stride, not for filtering the rows in the table. | + +## 示例 + +```sql +set geaflow.dsl.clickhouse.url = 'jdbc:clickhouse://localhost:8123/default'; +set geaflow.dsl.clickhouse.username = 'default'; +set geaflow.dsl.clickhouse.password = ''; + +CREATE TABLE clickhouse_source_table ( + id BIGINT, + name VARCHAR, + score DOUBLE +) WITH ( + type='clickhouse', + geaflow.dsl.clickhouse.table.name = 'source_table' +); + +CREATE TABLE clickhouse_sink_table ( + id BIGINT, + name VARCHAR, + score DOUBLE +) WITH ( + type='clickhouse', + geaflow.dsl.clickhouse.table.name = 'sink_table', + geaflow.dsl.clickhouse.write.batch.size = '2000' +); + +INSERT INTO clickhouse_sink_table +SELECT * FROM clickhouse_source_table; +``` diff --git a/docs/docs-cn/source/5.application-development/3.connector/index.rst b/docs/docs-cn/source/5.application-development/3.connector/index.rst index 5ca142767..41fbbc7da 100644 --- a/docs/docs-cn/source/5.application-development/3.connector/index.rst +++ b/docs/docs-cn/source/5.application-development/3.connector/index.rst @@ -16,4 +16,5 @@ 8.hudi.md 9.pulsar.md 10.udc.md + 11.clickhouse.md diff --git a/docs/docs-en/source/5.application-development/3.connector/11.clickhouse.md b/docs/docs-en/source/5.application-development/3.connector/11.clickhouse.md new file mode 100644 index 000000000..17fdc3756 --- /dev/null +++ b/docs/docs-en/source/5.application-development/3.connector/11.clickhouse.md @@ -0,0 +1,67 @@ +# ClickHouse Connector Introduction +The ClickHouse Connector supports both reading and writing operations. + +The sink buffers rows and flushes them to ClickHouse in bulk (batched inserts). Because ClickHouse is columnar and optimized for bulk ingestion, batched writes give far higher throughput than inserting row by row through the generic JDBC connector. The batch size is configurable. + +The source can split a read into several partitions over a numeric column so the partitions can be read in parallel. + +## Syntax + +```sql +CREATE TABLE clickhouse_table ( + id BIGINT, + name VARCHAR, + score DOUBLE +) WITH ( + type='clickhouse', + geaflow.dsl.clickhouse.url = 'jdbc:clickhouse://localhost:8123/default', + geaflow.dsl.clickhouse.username = 'default', + geaflow.dsl.clickhouse.password = '', + geaflow.dsl.clickhouse.table.name = 'source_table' +); +``` + +## Options + +| Key | Required | Description | +| -------- |------|---------------------------------------------------| +| geaflow.dsl.clickhouse.url | true | The ClickHouse JDBC url, e.g. jdbc:clickhouse://host:8123/database. | +| geaflow.dsl.clickhouse.table.name | true | The ClickHouse table name. | +| geaflow.dsl.clickhouse.username | false | The ClickHouse username, default 'default'. | +| geaflow.dsl.clickhouse.password | false | The ClickHouse password, default empty. | +| geaflow.dsl.clickhouse.write.batch.size | false | The number of rows buffered before a batch is flushed to ClickHouse. Larger batches give higher write throughput. Default 1000. | +| geaflow.dsl.clickhouse.driver | false | The ClickHouse JDBC driver class, default com.clickhouse.jdbc.ClickHouseDriver. | +| geaflow.dsl.clickhouse.partition.num | false | The number of source read partitions, default 1. | +| geaflow.dsl.clickhouse.partition.column | false | The column used to split the source read into partitions. Default value is 'id'. | +| geaflow.dsl.clickhouse.partition.lowerbound | false | The lowerbound used to decide the partition stride, not for filtering the rows in the table. | +| geaflow.dsl.clickhouse.partition.upperbound | false | The upperbound used to decide the partition stride, not for filtering the rows in the table. | + +## Example + +```sql +set geaflow.dsl.clickhouse.url = 'jdbc:clickhouse://localhost:8123/default'; +set geaflow.dsl.clickhouse.username = 'default'; +set geaflow.dsl.clickhouse.password = ''; + +CREATE TABLE clickhouse_source_table ( + id BIGINT, + name VARCHAR, + score DOUBLE +) WITH ( + type='clickhouse', + geaflow.dsl.clickhouse.table.name = 'source_table' +); + +CREATE TABLE clickhouse_sink_table ( + id BIGINT, + name VARCHAR, + score DOUBLE +) WITH ( + type='clickhouse', + geaflow.dsl.clickhouse.table.name = 'sink_table', + geaflow.dsl.clickhouse.write.batch.size = '2000' +); + +INSERT INTO clickhouse_sink_table +SELECT * FROM clickhouse_source_table; +``` diff --git a/docs/docs-en/source/5.application-development/3.connector/index.rst b/docs/docs-en/source/5.application-development/3.connector/index.rst index 7e696edc4..2b6e6bd14 100644 --- a/docs/docs-en/source/5.application-development/3.connector/index.rst +++ b/docs/docs-en/source/5.application-development/3.connector/index.rst @@ -15,4 +15,5 @@ Connector 7.hbase.md 8.hudi.md 9.pulsar.md - 10.udc.md \ No newline at end of file + 10.udc.md + 11.clickhouse.md \ No newline at end of file diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/pom.xml b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/pom.xml new file mode 100644 index 000000000..e5ec03cba --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/pom.xml @@ -0,0 +1,75 @@ + + + + + + org.apache.geaflow + geaflow-dsl-connector + 0.8.0-SNAPSHOT + + 4.0.0 + + geaflow-dsl-connector-clickhouse + geaflow-dsl-connector-clickhouse + + + 0.6.5 + 5.2.1 + 1.19.8 + + + + + org.apache.geaflow + geaflow-dsl-common + + + org.apache.geaflow + geaflow-dsl-connector-api + + + + com.clickhouse + clickhouse-jdbc + ${clickhouse-jdbc.version} + + + + org.apache.httpcomponents.client5 + httpclient5 + ${httpclient5.version} + + + + org.testng + testng + ${testng.version} + test + + + org.testcontainers + testcontainers + ${testcontainers.version} + test + + + diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseConfigKeys.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseConfigKeys.java new file mode 100644 index 000000000..950010fce --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseConfigKeys.java @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.clickhouse; + +import org.apache.geaflow.common.config.ConfigKey; +import org.apache.geaflow.common.config.ConfigKeys; + +public class ClickHouseConfigKeys { + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_DRIVER = ConfigKeys + .key("geaflow.dsl.clickhouse.driver") + .defaultValue("com.clickhouse.jdbc.ClickHouseDriver") + .description("The ClickHouse JDBC driver class, " + + "default com.clickhouse.jdbc.ClickHouseDriver."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_URL = ConfigKeys + .key("geaflow.dsl.clickhouse.url") + .noDefaultValue() + .description("The ClickHouse JDBC url, e.g. jdbc:clickhouse://host:8123/database."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_USERNAME = ConfigKeys + .key("geaflow.dsl.clickhouse.username") + .defaultValue("default") + .description("The ClickHouse username, default 'default'."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_PASSWORD = ConfigKeys + .key("geaflow.dsl.clickhouse.password") + .defaultValue("") + .description("The ClickHouse password, default empty."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_TABLE_NAME = ConfigKeys + .key("geaflow.dsl.clickhouse.table.name") + .noDefaultValue() + .description("The ClickHouse table name."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_WRITE_BATCH_SIZE = ConfigKeys + .key("geaflow.dsl.clickhouse.write.batch.size") + .defaultValue(1000) + .description("The number of rows buffered before a batch is flushed to ClickHouse. " + + "ClickHouse is columnar and strongly prefers bulk inserts, so larger batches " + + "give much higher write throughput. Default 1000."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_PARTITION_NUM = ConfigKeys + .key("geaflow.dsl.clickhouse.partition.num") + .defaultValue(1L) + .description("The number of source read partitions, default 1."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_PARTITION_COLUMN = ConfigKeys + .key("geaflow.dsl.clickhouse.partition.column") + .defaultValue("id") + .description("The column used to split the source read into partitions."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_PARTITION_LOWERBOUND = ConfigKeys + .key("geaflow.dsl.clickhouse.partition.lowerbound") + .defaultValue(0L) + .description("The lowerbound used to decide the partition stride, " + + "not for filtering the rows in the table."); + + public static final ConfigKey GEAFLOW_DSL_CLICKHOUSE_PARTITION_UPPERBOUND = ConfigKeys + .key("geaflow.dsl.clickhouse.partition.upperbound") + .defaultValue(0L) + .description("The upperbound used to decide the partition stride, " + + "not for filtering the rows in the table."); +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnector.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnector.java new file mode 100644 index 000000000..801b09d4f --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnector.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.clickhouse; + +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.dsl.connector.api.TableReadableConnector; +import org.apache.geaflow.dsl.connector.api.TableSink; +import org.apache.geaflow.dsl.connector.api.TableSource; +import org.apache.geaflow.dsl.connector.api.TableWritableConnector; + +public class ClickHouseTableConnector implements TableReadableConnector, TableWritableConnector { + + public static final String TYPE = "CLICKHOUSE"; + + @Override + public String getType() { + return TYPE; + } + + @Override + public TableSource createSource(Configuration conf) { + return new ClickHouseTableSource(); + } + + @Override + public TableSink createSink(Configuration conf) { + return new ClickHouseTableSink(); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSink.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSink.java new file mode 100644 index 000000000..72f293445 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSink.java @@ -0,0 +1,144 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.clickhouse; + +import java.io.IOException; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.util.List; +import org.apache.geaflow.api.context.RuntimeContext; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.StructType; +import org.apache.geaflow.dsl.common.types.TableField; +import org.apache.geaflow.dsl.connector.api.TableSink; +import org.apache.geaflow.dsl.connector.clickhouse.util.ClickHouseUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * ClickHouse table sink. Unlike a row-by-row JDBC INSERT, rows are buffered onto a reusable + * prepared statement and flushed in bulk once the configured batch size is reached (and again when + * the window finishes). ClickHouse is columnar and ingests far more efficiently from bulk inserts, + * so this is the main performance win over connecting through the generic JDBC connector. + */ +public class ClickHouseTableSink implements TableSink { + + private static final Logger LOGGER = LoggerFactory.getLogger(ClickHouseTableSink.class); + + private StructType schema; + private String driver; + private String url; + private String username; + private String password; + private String tableName; + private int batchSize; + + private transient Connection connection; + private transient PreparedStatement statement; + private transient int bufferedRows; + + @Override + public void init(Configuration tableConf, StructType tableSchema) { + LOGGER.info("init clickhouse sink with config: {}, schema: {}", tableConf, tableSchema); + this.schema = tableSchema; + this.driver = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_DRIVER); + this.url = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_URL); + this.username = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_USERNAME); + this.password = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PASSWORD); + this.tableName = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_TABLE_NAME); + this.batchSize = + tableConf.getInteger(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_WRITE_BATCH_SIZE); + if (this.batchSize <= 0) { + throw new GeaFlowDSLException("clickhouse write batch size must be > 0, but was {}", + this.batchSize); + } + } + + @Override + public void open(RuntimeContext context) { + try { + Class.forName(this.driver); + this.connection = DriverManager.getConnection(url, username, password); + List fields = schema.getFields(); + String insertSql = ClickHouseUtils.buildInsertSql(tableName, fields); + this.statement = connection.prepareStatement(insertSql); + this.bufferedRows = 0; + LOGGER.info("open clickhouse sink for table {} with batch size {}", tableName, batchSize); + } catch (ClassNotFoundException e) { + throw new GeaFlowDSLException("failed to load clickhouse driver: " + driver, e); + } catch (SQLException e) { + throw new GeaFlowDSLException("failed to connect to clickhouse: " + url, e); + } + } + + @Override + public void write(Row row) throws IOException { + try { + ClickHouseUtils.bindRow(statement, schema.getFields(), row); + statement.addBatch(); + bufferedRows++; + if (bufferedRows >= batchSize) { + flush(); + } + } catch (SQLException e) { + throw new GeaFlowDSLException("failed to buffer row for table: " + tableName, e); + } + } + + @Override + public void finish() throws IOException { + try { + flush(); + } catch (SQLException e) { + throw new GeaFlowDSLException("failed to flush batch to table: " + tableName, e); + } + } + + private void flush() throws SQLException { + if (bufferedRows == 0) { + return; + } + statement.executeBatch(); + statement.clearBatch(); + LOGGER.info("flushed {} rows to clickhouse table {}", bufferedRows, tableName); + bufferedRows = 0; + } + + @Override + public void close() { + try { + if (this.statement != null) { + this.statement.close(); + this.statement = null; + } + if (this.connection != null) { + this.connection.close(); + this.connection = null; + } + LOGGER.info("close clickhouse sink"); + } catch (SQLException e) { + throw new GeaFlowDSLException("failed to close clickhouse sink", e); + } + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSource.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSource.java new file mode 100644 index 000000000..42d7bf400 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableSource.java @@ -0,0 +1,280 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.clickhouse; + +import java.io.IOException; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import org.apache.geaflow.api.context.RuntimeContext; +import org.apache.geaflow.api.window.WindowType; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.StructType; +import org.apache.geaflow.dsl.common.types.TableSchema; +import org.apache.geaflow.dsl.connector.api.FetchData; +import org.apache.geaflow.dsl.connector.api.Offset; +import org.apache.geaflow.dsl.connector.api.Partition; +import org.apache.geaflow.dsl.connector.api.TableSource; +import org.apache.geaflow.dsl.connector.api.serde.DeserializerFactory; +import org.apache.geaflow.dsl.connector.api.serde.TableDeserializer; +import org.apache.geaflow.dsl.connector.api.window.FetchWindow; +import org.apache.geaflow.dsl.connector.clickhouse.util.ClickHouseUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * ClickHouse table source. The read is split into partitions over a numeric column so partitions + * can be fetched in parallel, each partition paging through bounded windows. + */ +public class ClickHouseTableSource implements TableSource { + + private static final Logger LOGGER = LoggerFactory.getLogger(ClickHouseTableSource.class); + + private StructType schema; + private String driver; + private String url; + private String username; + private String password; + private String tableName; + private long partitionNum; + private String partitionColumn; + private long lowerBound; + private long upperBound; + private final Map partitionConnectionMap = new HashMap<>(); + private final Map partitionStatementMap = new HashMap<>(); + + @Override + public void init(Configuration tableConf, TableSchema tableSchema) { + LOGGER.info("init clickhouse source with config: {}, schema: {}", tableConf, tableSchema); + this.schema = tableSchema; + this.driver = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_DRIVER); + this.url = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_URL); + this.username = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_USERNAME); + this.password = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PASSWORD); + this.tableName = tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_TABLE_NAME); + this.partitionNum = + tableConf.getLong(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_NUM); + if (this.partitionNum <= 0) { + throw new GeaFlowDSLException("invalid clickhouse partition number: {}", partitionNum); + } + this.partitionColumn = + tableConf.getString(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_COLUMN); + this.lowerBound = + tableConf.getLong(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_LOWERBOUND); + this.upperBound = + tableConf.getLong(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_UPPERBOUND); + if (partitionNum > 1) { + if (lowerBound >= upperBound) { + throw new GeaFlowDSLException("upperbound must be greater than lowerbound " + + "(lowerbound:%d upperbound:%d).", lowerBound, upperBound); + } + partitionNum = Math.min(upperBound - lowerBound, partitionNum); + } + } + + @Override + public void open(RuntimeContext context) { + try { + Class.forName(this.driver); + } catch (ClassNotFoundException e) { + throw new GeaFlowDSLException("failed to load clickhouse driver: " + driver, e); + } + } + + @Override + public List listPartitions() { + if (partitionNum <= 0) { + throw new GeaFlowDSLException("invalid clickhouse partition number: {}", partitionNum); + } + if (partitionNum == 1) { + return Collections.singletonList(new ClickHousePartition(tableName, "")); + } + long stride = upperBound / partitionNum - lowerBound / partitionNum; + long currentValue = lowerBound; + List partitions = new ArrayList<>(); + for (long i = 0; i < partitionNum; ++i) { + String lBound = i != 0 ? String.format("%s >= %d", partitionColumn, currentValue) : null; + currentValue += stride; + String uBound = i != partitionNum - 1 + ? String.format("%s < %d", partitionColumn, currentValue) : null; + String whereClause; + if (uBound == null) { + whereClause = lBound; + } else if (lBound == null) { + whereClause = String.format("%s OR %s IS NULL", uBound, partitionColumn); + } else { + whereClause = String.format("%s AND %s", lBound, uBound); + } + partitions.add(new ClickHousePartition(tableName, "WHERE " + whereClause)); + } + return partitions; + } + + @Override + public List listPartitions(int parallelism) { + return listPartitions(); + } + + @Override + public TableDeserializer getDeserializer(Configuration conf) { + return DeserializerFactory.loadRowTableDeserializer(); + } + + @Override + public FetchData fetch(Partition partition, Optional startOffset, + FetchWindow windowInfo) throws IOException { + if (!(windowInfo.getType() == WindowType.SIZE_TUMBLING_WINDOW + || windowInfo.getType() == WindowType.ALL_WINDOW)) { + throw new GeaFlowDSLException("not support window type: {}", windowInfo.getType()); + } + ClickHousePartition chPartition = (ClickHousePartition) partition; + if (!chPartition.getTableName().equals(this.tableName)) { + throw new GeaFlowDSLException("wrong partition"); + } + Statement statement = partitionStatementMap.get(partition); + if (statement == null) { + try { + Connection connection = DriverManager.getConnection(url, username, password); + statement = connection.createStatement(); + partitionConnectionMap.put(partition, connection); + partitionStatementMap.put(partition, statement); + } catch (SQLException e) { + throw new GeaFlowDSLException("failed to connect to clickhouse: " + url, e); + } + } + + long offset = 0; + if (startOffset.isPresent()) { + offset = startOffset.get().getOffset(); + } + List dataList; + try { + dataList = ClickHouseUtils.selectRowsFromTable(statement, this.tableName, + chPartition.getWhereClause(), this.schema.size(), offset, windowInfo.windowSize(), + this.schema.getField(0).getName()); + } catch (SQLException e) { + throw new GeaFlowDSLException("failed to select rows from table: " + tableName, e); + } + ClickHouseOffset nextOffset = new ClickHouseOffset(offset + dataList.size()); + boolean isFinish = windowInfo.getType() == WindowType.ALL_WINDOW + || dataList.size() < windowInfo.windowSize(); + return (FetchData) FetchData.createStreamFetch(dataList, nextOffset, isFinish); + } + + @Override + public void close() { + try { + for (Statement statement : this.partitionStatementMap.values()) { + if (statement != null) { + statement.close(); + } + } + this.partitionStatementMap.clear(); + for (Connection connection : this.partitionConnectionMap.values()) { + if (connection != null) { + connection.close(); + } + } + this.partitionConnectionMap.clear(); + } catch (SQLException e) { + throw new GeaFlowDSLException("failed to close clickhouse connection.", e); + } + } + + public static class ClickHousePartition implements Partition { + + private final String tableName; + private final String whereClause; + + public ClickHousePartition(String tableName, String whereClause) { + this.tableName = tableName; + this.whereClause = whereClause; + } + + public String getTableName() { + return tableName; + } + + public String getWhereClause() { + return whereClause; + } + + @Override + public String getName() { + if (whereClause == null || whereClause.isEmpty()) { + return tableName; + } else { + return tableName + "-" + whereClause; + } + } + + @Override + public int hashCode() { + return Objects.hash(tableName, whereClause); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ClickHousePartition)) { + return false; + } + ClickHousePartition that = (ClickHousePartition) o; + return Objects.equals(tableName, that.tableName) + && Objects.equals(whereClause, that.whereClause); + } + } + + public static class ClickHouseOffset implements Offset { + + private final long offset; + + public ClickHouseOffset(long offset) { + this.offset = offset; + } + + @Override + public String humanReadable() { + return String.valueOf(offset); + } + + @Override + public long getOffset() { + return offset; + } + + @Override + public boolean isTimestamp() { + return false; + } + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/util/ClickHouseUtils.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/util/ClickHouseUtils.java new file mode 100644 index 000000000..1799a2ab3 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/java/org/apache/geaflow/dsl/connector/clickhouse/util/ClickHouseUtils.java @@ -0,0 +1,111 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.clickhouse.util; + +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.geaflow.common.type.IType; +import org.apache.geaflow.common.type.Types; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.impl.ObjectRow; +import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; +import org.apache.geaflow.dsl.common.types.TableField; +import org.apache.geaflow.dsl.common.util.Windows; + +public class ClickHouseUtils { + + /** + * Builds a parameterized batch-insert statement of the form + * {@code INSERT INTO table (c1, c2, ...) VALUES (?, ?, ...)}. The same prepared statement is + * reused across the whole batch, which is what lets ClickHouse ingest rows in bulk instead of + * one round-trip per row. + */ + public static String buildInsertSql(String tableName, List fields) { + String columns = fields.stream() + .map(TableField::getName) + .collect(Collectors.joining(", ")); + String placeholders = fields.stream() + .map(field -> "?") + .collect(Collectors.joining(", ")); + return String.format("INSERT INTO %s (%s) VALUES (%s)", tableName, columns, placeholders); + } + + /** + * Binds a single row onto the prepared statement parameters, mapping GeaFlow types to JDBC + * setters. The caller is expected to call {@link PreparedStatement#addBatch()} afterwards. + */ + public static void bindRow(PreparedStatement statement, List fields, Row row) + throws SQLException { + for (int i = 0; i < fields.size(); i++) { + TableField field = fields.get(i); + IType type = field.getType(); + Object value = row.getField(i, type); + int paramIndex = i + 1; + if (value == null) { + if (!field.isNullable()) { + throw new GeaFlowDSLException("field " + field.getName() + " can not be null"); + } + statement.setObject(paramIndex, null); + continue; + } + switch (type.getName()) { + case Types.TYPE_NAME_STRING: + case Types.TYPE_NAME_BINARY_STRING: + statement.setString(paramIndex, value.toString()); + break; + default: + statement.setObject(paramIndex, value); + } + } + } + + /** + * Reads a window of rows from the table, optionally restricted to a partition's where clause. + * Mirrors the JDBC connector's paging so a partition can be fetched in bounded windows. + */ + public static List selectRowsFromTable(Statement statement, String tableName, + String whereClause, int columnNum, long startOffset, + long windowSize, String orderByColumnName) + throws SQLException { + if (windowSize == Windows.SIZE_OF_ALL_WINDOW) { + windowSize = Integer.MAX_VALUE; + } else if (windowSize <= 0) { + throw new GeaFlowDSLException("wrong windowSize"); + } + String selectQuery = String.format("SELECT * FROM %s %s ORDER BY %s LIMIT %s OFFSET %s", + tableName, whereClause, orderByColumnName, windowSize, startOffset); + ResultSet resultSet = statement.executeQuery(selectQuery); + List rowList = new ArrayList<>(); + while (resultSet.next()) { + Object[] values = new Object[columnNum]; + for (int i = 1; i <= columnNum; i++) { + values[i - 1] = resultSet.getObject(i); + } + rowList.add(ObjectRow.create(values)); + } + resultSet.close(); + return rowList; + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/resources/META-INF/services/org.apache.geaflow.dsl.connector.api.TableConnector b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/resources/META-INF/services/org.apache.geaflow.dsl.connector.api.TableConnector new file mode 100644 index 000000000..8c422bb96 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/main/resources/META-INF/services/org.apache.geaflow.dsl.connector.api.TableConnector @@ -0,0 +1,20 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +org.apache.geaflow.dsl.connector.clickhouse.ClickHouseTableConnector diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/test/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnectorTest.java b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/test/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnectorTest.java new file mode 100644 index 000000000..b2c1530a3 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/geaflow-dsl-connector-clickhouse/src/test/java/org/apache/geaflow/dsl/connector/clickhouse/ClickHouseTableConnectorTest.java @@ -0,0 +1,280 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.dsl.connector.clickhouse; + +import java.io.IOException; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Locale; +import java.util.Optional; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.common.type.Types; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.impl.ObjectRow; +import org.apache.geaflow.dsl.common.types.StructType; +import org.apache.geaflow.dsl.common.types.TableField; +import org.apache.geaflow.dsl.common.types.TableSchema; +import org.apache.geaflow.dsl.connector.api.FetchData; +import org.apache.geaflow.dsl.connector.api.Partition; +import org.apache.geaflow.dsl.connector.api.TableConnector; +import org.apache.geaflow.dsl.connector.api.util.ConnectorFactory; +import org.apache.geaflow.dsl.connector.api.window.AllFetchWindow; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.utility.DockerImageName; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +public class ClickHouseTableConnectorTest { + + private static final Logger LOGGER = + LoggerFactory.getLogger(ClickHouseTableConnectorTest.class); + + private static final DockerImageName IMAGE = + DockerImageName.parse("clickhouse/clickhouse-server:24.3"); + + private static final int ROWS = 500; + + private static GenericContainer clickhouse; + private static String url; + private static String username; + private static String password; + + @BeforeClass + public static void setup() throws SQLException { + username = "ch_user"; + password = "ch_pass"; + clickhouse = new GenericContainer<>(IMAGE); + clickhouse.withExposedPorts(8123); + clickhouse.withEnv("CLICKHOUSE_USER", username); + clickhouse.withEnv("CLICKHOUSE_PASSWORD", password); + clickhouse.withEnv("CLICKHOUSE_DB", "default"); + clickhouse.withEnv("CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT", "1"); + clickhouse.waitingFor(Wait.forHttp("/ping").forStatusCode(200)); + clickhouse.start(); + url = "jdbc:clickhouse://" + clickhouse.getHost() + ":" + + clickhouse.getMappedPort(8123) + "/default"; + LOGGER.info("clickhouse started at {}", url); + + try (Connection connection = DriverManager.getConnection(url, username, password); + Statement statement = connection.createStatement()) { + statement.execute("CREATE TABLE events (id Int64, name String, score Float64) " + + "ENGINE = MergeTree() ORDER BY id"); + // Seed the read tests with a known, ordered data set. + StringBuilder insert = new StringBuilder("INSERT INTO events (id, name, score) VALUES "); + for (int i = 0; i < ROWS; i++) { + if (i > 0) { + insert.append(","); + } + insert.append(String.format(Locale.ROOT, "(%d, 'name-%d', %f)", i, i, (double) i)); + } + statement.execute(insert.toString()); + } + } + + @AfterClass + public static void cleanup() { + if (clickhouse != null) { + clickhouse.stop(); + } + } + + private static TableField[] fields() { + return new TableField[]{ + new TableField("id", Types.LONG, false), + new TableField("name", Types.BINARY_STRING, false), + new TableField("score", Types.DOUBLE, false) + }; + } + + private static Configuration baseConf(String tableName) { + Configuration conf = new Configuration(); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_URL, url); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_USERNAME, username); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PASSWORD, password); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_TABLE_NAME, tableName); + return conf; + } + + private static long countRows(String tableName) throws SQLException { + try (Connection connection = DriverManager.getConnection(url, username, password); + Statement statement = connection.createStatement(); + ResultSet rs = statement.executeQuery("SELECT count() FROM " + tableName)) { + rs.next(); + return rs.getLong(1); + } + } + + private static void createTable(String tableName, String engine) throws SQLException { + try (Connection connection = DriverManager.getConnection(url, username, password); + Statement statement = connection.createStatement()) { + statement.execute(String.format("CREATE TABLE %s (id Int64, name String, score Float64) " + + "ENGINE = %s", tableName, engine)); + } + } + + private static void writeRows(String tableName, int batchSize, int count) throws IOException { + Configuration conf = baseConf(tableName); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_WRITE_BATCH_SIZE, + String.valueOf(batchSize)); + ClickHouseTableSink sink = new ClickHouseTableSink(); + sink.init(conf, new StructType(fields())); + sink.open(null); + try { + for (int i = 0; i < count; i++) { + sink.write(ObjectRow.create(new Object[]{(long) i, "name-" + i, (double) i})); + } + sink.finish(); + } finally { + sink.close(); + } + } + + private static List readAll(Configuration conf) throws IOException { + ClickHouseTableSource source = new ClickHouseTableSource(); + source.init(conf, new TableSchema(fields())); + source.open(null); + List rows = new ArrayList<>(); + try { + for (Partition partition : source.listPartitions()) { + FetchData data = source.fetch(partition, Optional.empty(), + new AllFetchWindow(0L)); + Iterator iterator = data.getDataIterator(); + while (iterator.hasNext()) { + rows.add(iterator.next()); + } + } + } finally { + source.close(); + } + return rows; + } + + @Test + public void testSinkBatchWrite() throws Exception { + createTable("sink_test", "MergeTree() ORDER BY id"); + writeRows("sink_test", 100, ROWS); + Assert.assertEquals(countRows("sink_test"), ROWS); + } + + @Test + public void testSourceReadAll() throws Exception { + Configuration conf = baseConf("events"); + List rows = readAll(conf); + Assert.assertEquals(rows.size(), ROWS); + // Rows come back ordered by id; spot check that the full first row round-tripped. + Row first = rows.get(0); + Assert.assertEquals(((Number) first.getField(0, Types.LONG)).longValue(), 0L); + Assert.assertEquals(String.valueOf(first.getField(1, Types.BINARY_STRING)), "name-0"); + Assert.assertEquals(((Number) first.getField(2, Types.DOUBLE)).doubleValue(), 0.0, 1e-9); + } + + @Test + public void testSourceParallelPartitionedRead() throws Exception { + Configuration conf = baseConf("events"); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_NUM, "2"); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_COLUMN, "id"); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_LOWERBOUND, "0"); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_UPPERBOUND, + String.valueOf(ROWS)); + + ClickHouseTableSource source = new ClickHouseTableSource(); + source.init(conf, new TableSchema(fields())); + source.open(null); + List partitions = source.listPartitions(); + Assert.assertEquals(partitions.size(), 2); + + int total = 0; + for (Partition partition : partitions) { + FetchData data = source.fetch(partition, Optional.empty(), new AllFetchWindow(0L)); + Iterator iterator = data.getDataIterator(); + while (iterator.hasNext()) { + iterator.next(); + total++; + } + } + source.close(); + // The two partitions cover disjoint id ranges; together they must return every row exactly once. + Assert.assertEquals(total, ROWS); + } + + @Test + public void testListPartitionsSplitsRange() { + Configuration conf = baseConf("events"); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_NUM, "4"); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_LOWERBOUND, "0"); + conf.put(ClickHouseConfigKeys.GEAFLOW_DSL_CLICKHOUSE_PARTITION_UPPERBOUND, "400"); + ClickHouseTableSource source = new ClickHouseTableSource(); + source.init(conf, new TableSchema(fields())); + List partitions = source.listPartitions(); + Assert.assertEquals(partitions.size(), 4); + for (Partition partition : partitions) { + Assert.assertTrue(partition.getName().contains("WHERE")); + } + } + + @Test + public void testConnectorDiscoveredViaSpi() { + // The DSL resolves `type='clickhouse'` through ConnectorFactory + the ServiceLoader SPI; + // this verifies the META-INF/services registration is actually picked up. + TableConnector connector = ConnectorFactory.loadConnector("clickhouse"); + Assert.assertTrue(connector instanceof ClickHouseTableConnector); + Assert.assertEquals(connector.getType(), ClickHouseTableConnector.TYPE); + } + + @Test + public void testBatchWriteOutperformsRowByRow() throws Exception { + createTable("bench_batched", "Memory"); + createTable("bench_rowbyrow", "Memory"); + int benchRows = 2000; + + long batchedStart = System.nanoTime(); + writeRows("bench_batched", 2000, benchRows); + long batchedNanos = System.nanoTime() - batchedStart; + + // batch size 1 flushes every row, i.e. one insert round-trip per row (the row-by-row baseline). + long rowByRowStart = System.nanoTime(); + writeRows("bench_rowbyrow", 1, benchRows); + long rowByRowNanos = System.nanoTime() - rowByRowStart; + + Assert.assertEquals(countRows("bench_batched"), benchRows); + Assert.assertEquals(countRows("bench_rowbyrow"), benchRows); + + double batchedPerSec = benchRows / (batchedNanos / 1_000_000_000.0); + double rowByRowPerSec = benchRows / (rowByRowNanos / 1_000_000_000.0); + LOGGER.info("clickhouse write benchmark for {} rows: batched={} rows/s, " + + "row-by-row={} rows/s, speedup={}x", + benchRows, String.format("%.0f", batchedPerSec), String.format("%.0f", rowByRowPerSec), + String.format("%.1f", batchedPerSec / rowByRowPerSec)); + + Assert.assertTrue(batchedNanos < rowByRowNanos, + "batched write should be faster than row-by-row"); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml b/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml index e1f262b9d..f35358649 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml +++ b/geaflow/geaflow-dsl/geaflow-dsl-connector/pom.xml @@ -49,6 +49,7 @@ geaflow-dsl-connector-paimon geaflow-dsl-connector-neo4j geaflow-dsl-connector-elasticsearch + geaflow-dsl-connector-clickhouse