Skip to content

Commit 889c0a6

Browse files
committed
Add predicate pushdown for equality on Clickhouse map values
1 parent 2d65996 commit 889c0a6

File tree

2 files changed

+55
-14
lines changed

2 files changed

+55
-14
lines changed

plugin/trino-clickhouse/src/main/java/io/trino/plugin/clickhouse/ClickHouseClient.java

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import io.trino.plugin.jdbc.LongWriteFunction;
4949
import io.trino.plugin.jdbc.ObjectReadFunction;
5050
import io.trino.plugin.jdbc.ObjectWriteFunction;
51+
import io.trino.plugin.jdbc.PreparedQuery;
5152
import io.trino.plugin.jdbc.QueryBuilder;
5253
import io.trino.plugin.jdbc.RemoteTableName;
5354
import io.trino.plugin.jdbc.SliceWriteFunction;
@@ -74,6 +75,8 @@
7475
import io.trino.spi.connector.ColumnPosition;
7576
import io.trino.spi.connector.ConnectorSession;
7677
import io.trino.spi.connector.ConnectorTableMetadata;
78+
import io.trino.spi.connector.JoinStatistics;
79+
import io.trino.spi.connector.JoinType;
7780
import io.trino.spi.expression.ConnectorExpression;
7881
import io.trino.spi.expression.Variable;
7982
import io.trino.spi.type.CharType;
@@ -251,6 +254,13 @@ public ClickHouseClient(
251254
.map("$not($is_null(value))").to("value IS NOT NULL")
252255
.map("$not(value: boolean)").to("NOT value")
253256
.map("$is_null(value)").to("value IS NULL")
257+
.map("$equal(left, right)").to("left = right")
258+
.map("$not_equal(left, right)").to("left <> right")
259+
.map("$less_than(left, right)").to("left < right")
260+
.map("$less_than_or_equal(left, right)").to("left <= right")
261+
.map("$greater_than(left, right)").to("left > right")
262+
.map("$greater_than_or_equal(left, right)").to("left >= right")
263+
.map("$operator$subscript(map, value)").to("map[value]")
254264
.build();
255265
this.aggregateFunctionRewriter = new AggregateFunctionRewriter<>(
256266
this.connectorExpressionRewriter,
@@ -278,9 +288,36 @@ public Optional<JdbcExpression> implementAggregation(ConnectorSession session, A
278288
@Override
279289
public Optional<ParameterizedExpression> convertPredicate(ConnectorSession session, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
280290
{
291+
for (ColumnHandle columnHandle : assignments.values()) {
292+
JdbcColumnHandle jdbcColumnHandle = (JdbcColumnHandle) columnHandle;
293+
JdbcTypeHandle typeHandle = jdbcColumnHandle.getJdbcTypeHandle();
294+
String jdbcTypeName = typeHandle.jdbcTypeName()
295+
.orElseThrow(() -> new TrinoException(JDBC_ERROR, "Type name is missing: " + typeHandle));
296+
if (jdbcColumnHandle.getColumnType() instanceof VarcharType &&
297+
getUnsupportedTypeHandling(session) == CONVERT_TO_VARCHAR &&
298+
toColumnMapping(session, jdbcTypeName, typeHandle.jdbcType(), typeHandle.decimalDigits(), typeHandle.columnSize()).isEmpty()) {
299+
// Column is mapped to VARCHAR using unsupported type handling, predicate pushdown may not work properly
300+
return Optional.empty();
301+
}
302+
}
303+
281304
return connectorExpressionRewriter.rewrite(session, expression, assignments);
282305
}
283306

307+
@Override
308+
public Optional<PreparedQuery> implementJoin(
309+
ConnectorSession session,
310+
JoinType joinType,
311+
PreparedQuery leftSource,
312+
Map<JdbcColumnHandle, String> leftProjections,
313+
PreparedQuery rightSource,
314+
Map<JdbcColumnHandle, String> rightProjections,
315+
List<ParameterizedExpression> joinConditions,
316+
JoinStatistics statistics)
317+
{
318+
return Optional.empty();
319+
}
320+
284321
@Override
285322
public boolean supportsTopN(ConnectorSession session, JdbcTableHandle handle, List<JdbcSortItem> sortOrder)
286323
{

plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorTest.java

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -669,17 +669,21 @@ public void testInsertIntoNotNullColumn()
669669
@Override
670670
public void testInsertMap()
671671
{
672-
// TODO: Add more types here
673-
testMapRoundTrip("INTEGER", "2");
674-
testMapRoundTrip("VARCHAR", "CAST('foobar' AS VARCHAR)");
672+
try (TestTable table = newTrinoTable("test_insert_map_", "(col map(INTEGER, BIGINT) NOT NULL)")) {
673+
assertUpdate("INSERT INTO " + table.getName() + " VALUES map(ARRAY[1], ARRAY[BIGINT '123456789'])", 1);
674+
assertThat(query("SELECT col[1] FROM " + table.getName()))
675+
.matches("VALUES BIGINT '123456789'");
676+
}
675677
}
676678

677-
private void testMapRoundTrip(String valueType, String value)
679+
@Test
680+
public void testMapPredicatePushdown()
678681
{
679-
try (TestTable table = newTrinoTable("test_insert_map_", "(col map(INTEGER, %s) NOT NULL)".formatted(valueType))) {
680-
assertUpdate("INSERT INTO " + table.getName() + " VALUES map(ARRAY[1], ARRAY[%s])".formatted(value), 1);
681-
assertThat(query("SELECT col[1] FROM " + table.getName()))
682-
.matches("VALUES " + value);
682+
try (TestTable table = newTrinoTable("test_map_predicate_pushdown", "(id INT, map_t map(INTEGER, BIGINT) NOT NULL)")) {
683+
assertUpdate("INSERT INTO " + table.getName() + " VALUES (1, map(ARRAY[1], ARRAY[BIGINT '123456789']))", 1);
684+
assertThat(query("SELECT id FROM " + table.getName() + " WHERE map_t[1] = BIGINT '123456789'"))
685+
.matches("VALUES 1")
686+
.isFullyPushedDown();
683687
}
684688
}
685689

@@ -969,19 +973,19 @@ a_enum_2 Enum('hello', 'world', 'a', 'b', 'c', '%', '_'))
969973

970974
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_string = a_string_alias")).isFullyPushedDown();
971975
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_string = a_string_alias" + withConnectorExpression)).isFullyPushedDown();
972-
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_string = a_enum_1")).isNotFullyPushedDown(FilterNode.class);
973-
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_string = a_enum_1" + withConnectorExpression)).isNotFullyPushedDown(FilterNode.class);
976+
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_string = a_enum_1")).isFullyPushedDown();
977+
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_string = a_enum_1" + withConnectorExpression)).isFullyPushedDown();
974978
assertThat(query(convertToVarchar, "SELECT some_column FROM " + table.getName() + " WHERE a_string = unsupported_1")).isNotFullyPushedDown(FilterNode.class);
975979
assertThat(query(convertToVarchar, "SELECT some_column FROM " + table.getName() + " WHERE a_string = unsupported_1" + withConnectorExpression)).isNotFullyPushedDown(FilterNode.class);
976980

977981
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = 'hello'")).isNotFullyPushedDown(FilterNode.class);
978-
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = 'hello'" + withConnectorExpression)).isNotFullyPushedDown(FilterNode.class);
982+
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = 'hello'" + withConnectorExpression)).isFullyPushedDown();
979983
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = 'not_a_value'")).isNotFullyPushedDown(FilterNode.class);
980-
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = 'not_a_value'" + withConnectorExpression)).isNotFullyPushedDown(FilterNode.class);
984+
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = 'not_a_value'" + withConnectorExpression)).isFullyPushedDown();
981985
// pushdown of a condition, both sides of the same native type, which is mapped to varchar,
982986
// not allowed because some operations (e.g. inequalities) may not be allowed in the native system on an unknown native types
983-
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = a_enum_2")).isNotFullyPushedDown(FilterNode.class);
984-
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = a_enum_2" + withConnectorExpression)).isNotFullyPushedDown(FilterNode.class);
987+
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = a_enum_2")).isFullyPushedDown();
988+
assertThat(query("SELECT some_column FROM " + table.getName() + " WHERE a_enum_1 = a_enum_2" + withConnectorExpression)).isFullyPushedDown();
985989

986990
// pushdown of a condition, both sides of the same native type, which is mapped to varchar,
987991
// not allowed because some operations (e.g. inequalities) may not be allowed in the native system on an unknown native types

0 commit comments

Comments
 (0)