Skip to content

Commit 0d34e01

Browse files
azatazat
authored
Merge pull request ClickHouse#84515 from korowa/global-in-kv-get-keys
Support type casts for IN and GLOBAL IN over KV storages
2 parents 51f0b57 + 5b7e4c5 commit 0d34e01

File tree

5 files changed

+494
-45
lines changed

5 files changed

+494
-45
lines changed

‎src/Storages/KVStorageUtils.cpp‎

Lines changed: 83 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
#include <Storages/KVStorageUtils.h>
22

3+
#include <Columns/ColumnNullable.h>
34
#include <Columns/ColumnSet.h>
45
#include <Columns/ColumnConst.h>
6+
#include <Common/assert_cast.h>
7+
#include <DataTypes/Utils.h>
58

69
#include <Parsers/ASTIdentifier.h>
710
#include <Parsers/ASTLiteral.h>
811
#include <Parsers/ASTSelectQuery.h>
912

1013
#include <Interpreters/Set.h>
14+
#include <Interpreters/castColumn.h>
1115
#include <Interpreters/convertFieldToType.h>
1216
#include <Interpreters/evaluateConstantExpression.h>
1317

@@ -52,55 +56,34 @@ bool traverseDAGFilter(
5256
return false;
5357
return true;
5458
}
55-
if (func_name == "equals" || func_name == "in")
59+
if (func_name == "equals")
5660
{
5761
if (elem->children.size() != 2)
5862
return false;
5963

60-
if (func_name == "in")
61-
{
62-
const auto * key = elem->children.at(0);
63-
while (key->type == ActionsDAG::ActionType::ALIAS)
64-
key = key->children.at(0);
65-
66-
if (key->type != ActionsDAG::ActionType::INPUT)
67-
return false;
68-
69-
if (key->result_name != primary_key)
70-
return false;
71-
72-
const auto * value = elem->children.at(1);
73-
if (value->type != ActionsDAG::ActionType::COLUMN)
74-
return false;
75-
76-
const IColumn * value_col = value->column.get();
77-
if (const auto * col_const = typeid_cast<const ColumnConst *>(value_col))
78-
value_col = &col_const->getDataColumn();
79-
80-
const auto * col_set = typeid_cast<const ColumnSet *>(value_col);
81-
if (!col_set)
82-
return false;
83-
84-
auto future_set = col_set->getData();
85-
future_set->buildOrderedSetInplace(context);
86-
87-
auto set = future_set->get();
88-
if (!set)
89-
return false;
64+
const auto * key = elem->children.at(0);
65+
while (key->type == ActionsDAG::ActionType::ALIAS)
66+
key = key->children.at(0);
9067

91-
if (!set->hasExplicitSetElements())
92-
return false;
68+
if (key->type != ActionsDAG::ActionType::INPUT)
69+
return false;
9370

94-
set->checkColumnsNumber(1);
95-
const auto & set_column = *set->getSetElements()[0];
71+
if (key->result_name != primary_key)
72+
return false;
9673

97-
if (set_column.getDataType() != primary_key_type->getTypeId())
98-
return false;
74+
const auto * value = elem->children.at(1);
75+
if (value->type != ActionsDAG::ActionType::COLUMN)
76+
return false;
9977

100-
for (size_t row = 0; row < set_column.size(); ++row)
101-
res->push_back(set_column[row]);
102-
return true;
103-
}
78+
auto converted_field = convertFieldToType((*value->column)[0], *primary_key_type);
79+
if (!converted_field.isNull())
80+
res->push_back(converted_field);
81+
return true;
82+
}
83+
if (func_name == "in" || func_name == "globalIn")
84+
{
85+
if (elem->children.size() != 2)
86+
return false;
10487

10588
const auto * key = elem->children.at(0);
10689
while (key->type == ActionsDAG::ActionType::ALIAS)
@@ -116,10 +99,65 @@ bool traverseDAGFilter(
11699
if (value->type != ActionsDAG::ActionType::COLUMN)
117100
return false;
118101

119-
auto converted_field = convertFieldToType((*value->column)[0], *primary_key_type);
120-
if (!converted_field.isNull())
121-
res->push_back(converted_field);
122-
return true;
102+
const IColumn * value_col = value->column.get();
103+
if (const auto * col_const = typeid_cast<const ColumnConst *>(value_col))
104+
value_col = &col_const->getDataColumn();
105+
106+
const auto * col_set = typeid_cast<const ColumnSet *>(value_col);
107+
if (!col_set)
108+
return false;
109+
110+
auto future_set = col_set->getData();
111+
future_set->buildOrderedSetInplace(context);
112+
113+
auto set = future_set->get();
114+
if (!set)
115+
return false;
116+
117+
if (!set->hasExplicitSetElements())
118+
return false;
119+
120+
set->checkColumnsNumber(1);
121+
const auto set_column_ptr = set->getSetElements()[0];
122+
const auto set_data_type_ptr = set->getElementsTypes()[0];
123+
const ColumnWithTypeAndName set_column = {set_column_ptr, set_data_type_ptr, ""};
124+
125+
// In case set values can be safely casted to PK data type, simply convert them,
126+
// and return as filter.
127+
//
128+
// When safe cast is not guaranteed, we still can try to convert set values, and
129+
// push them to set filter if the whole set has been casted successfully. If some
130+
// of the set values haven't been casted, the set is considered as non-containing
131+
// keys.
132+
//
133+
// NOTE: this behavior may affect a queries like `key in (NULL)`, but right now it
134+
// already triggers full scan.
135+
if (canBeSafelyCast(set_data_type_ptr, primary_key_type))
136+
{
137+
const auto casted_set_ptr = castColumnAccurate(set_column, primary_key_type);
138+
const auto & casted_set = *casted_set_ptr;
139+
for (size_t row = 0; row < set_column_ptr->size(); ++row)
140+
res->push_back(casted_set[row]);
141+
return true;
142+
}
143+
else
144+
{
145+
const auto casted_set_ptr = castColumnAccurateOrNull(set_column, primary_key_type);
146+
const auto & casted_set_nullable = assert_cast<const ColumnNullable &>(*casted_set_ptr);
147+
const auto & casted_set_null_map = casted_set_nullable.getNullMapData();
148+
for (char8_t i : casted_set_null_map)
149+
{
150+
if (i != 0)
151+
return false;
152+
}
153+
154+
const auto & casted_set_inner = casted_set_nullable.getNestedColumn();
155+
for (size_t row = 0; row < casted_set_inner.size(); row++)
156+
{
157+
res->push_back(casted_set_inner[row]);
158+
}
159+
return true;
160+
}
123161
}
124162
return false;
125163
}

‎tests/queries/0_stateless/03578_distributed_kv_global_in.reference‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
-- RocksDB: set
2+
0 val-0
3+
0 val-0
4+
1 val-1
5+
1 val-1
6+
2 val-2
7+
2 val-2
8+
-- RocksDB: subquery
9+
0 val-0
10+
0 val-0
11+
1 val-1
12+
1 val-1
13+
2 val-2
14+
2 val-2
15+
-- Rows read:
16+
6
17+
15
18+
-- KeeperMap: set
19+
0 val-0
20+
0 val-0
21+
1 val-1
22+
1 val-1
23+
2 val-2
24+
2 val-2
25+
-- KeeperMap: subquery
26+
0 val-0
27+
0 val-0
28+
1 val-1
29+
1 val-1
30+
2 val-2
31+
2 val-2
32+
-- Rows read:
33+
6
34+
15

‎tests/queries/0_stateless/03578_distributed_kv_global_in.sql‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
-- Tags: no-fasttest
2+
3+
SET prefer_localhost_replica = 0;
4+
5+
DROP TABLE IF EXISTS 03578_rocksdb_local, 03578_rocksdb_dist;
6+
7+
CREATE TABLE IF NOT EXISTS 03578_rocksdb_local
8+
(
9+
key UInt64,
10+
val String
11+
)
12+
ENGINE = EmbeddedRocksDB()
13+
PRIMARY KEY key;
14+
15+
CREATE TABLE IF NOT EXISTS 03578_rocksdb_dist
16+
(
17+
key UInt64,
18+
val String
19+
)
20+
ENGINE = Distributed(test_cluster_two_shards_localhost, currentDatabase(), 03578_rocksdb_local);
21+
22+
INSERT INTO 03578_rocksdb_local SELECT number, 'val-' || number FROM numbers(1000);
23+
24+
SELECT '-- RocksDB: set';
25+
26+
SELECT *
27+
FROM 03578_rocksdb_dist
28+
WHERE key GLOBAL IN (0, 1, 2)
29+
ORDER BY 1, 2;
30+
31+
SELECT '-- RocksDB: subquery';
32+
33+
SELECT *
34+
FROM 03578_rocksdb_dist
35+
WHERE key GLOBAL IN (
36+
SELECT number FROM numbers(3)
37+
)
38+
ORDER BY 1, 2;
39+
40+
SYSTEM FLUSH LOGS query_log;
41+
42+
SELECT '-- Rows read:';
43+
44+
SELECT read_rows
45+
FROM system.query_log
46+
WHERE current_database = currentDatabase()
47+
AND type = 'QueryFinish'
48+
AND query LIKE '%FROM 03578_rocksdb_dist%'
49+
AND is_initial_query
50+
ORDER BY event_time_microseconds;
51+
52+
DROP TABLE 03578_rocksdb_local, 03578_rocksdb_dist;
53+
54+
DROP TABLE IF EXISTS 03578_keepermap_local, 03578_keepermap_dist;
55+
56+
CREATE TABLE IF NOT EXISTS 03578_keepermap_local
57+
(
58+
key UInt64,
59+
val String
60+
)
61+
ENGINE = KeeperMap('/' || currentDatabase() || '/test_03578_global_in')
62+
PRIMARY KEY (key);
63+
64+
CREATE TABLE IF NOT EXISTS 03578_keepermap_dist
65+
(
66+
key UInt64,
67+
val String
68+
)
69+
ENGINE = Distributed(test_cluster_two_shards_localhost, currentDatabase(), 03578_keepermap_local);
70+
71+
INSERT INTO 03578_keepermap_local SELECT number, 'val-' || number FROM numbers(1000);
72+
73+
SELECT '-- KeeperMap: set';
74+
75+
SELECT *
76+
FROM 03578_keepermap_dist
77+
WHERE key GLOBAL IN (0, 1, 2)
78+
ORDER BY 1, 2;
79+
80+
SELECT '-- KeeperMap: subquery';
81+
82+
SELECT *
83+
FROM 03578_keepermap_dist
84+
WHERE key GLOBAL IN (
85+
SELECT number FROM numbers(3)
86+
)
87+
ORDER BY 1, 2;
88+
89+
SYSTEM FLUSH LOGS query_log;
90+
91+
SELECT '-- Rows read:';
92+
93+
SELECT read_rows
94+
FROM system.query_log
95+
WHERE current_database = currentDatabase()
96+
AND type = 'QueryFinish'
97+
AND query LIKE '%FROM 03578_keepermap_dist%'
98+
AND is_initial_query
99+
ORDER BY event_time_microseconds;
100+
101+
DROP TABLE IF EXISTS 03578_keepermap_local, 03578_keepermap_dist;

‎tests/queries/0_stateless/03578_kv_in_type_casts.reference‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
-- RocksDB: set safe to cast
2+
0 val-0
3+
1 val-1
4+
-- RocksDB: set not safe to cast
5+
0 val-0
6+
1 val-1
7+
-- RocksDB: set with overflow
8+
0 val-0
9+
-- RocksDB: set with cast failure
10+
-- RocksDB: subquery safe to cast
11+
0 val-0
12+
1 val-1
13+
-- RocksDB: subquery not safe to cast
14+
0 val-0
15+
1 val-1
16+
-- RocksDB: subquery with overflow
17+
0 val-0
18+
-- RocksDB: subquery with cast failure
19+
0 val-0
20+
-- Rows read:
21+
2
22+
2
23+
1
24+
4
25+
4
26+
102
27+
102
28+
29+
-- RocksDB null: set without null
30+
0 val-0
31+
1 val-1
32+
-- RocksDB null: set with null
33+
1 val-1
34+
\N val-null
35+
-- RocksDB null: subquery without null
36+
0 val-0
37+
1 val-1
38+
-- RocksDB null: subquery with null
39+
1 val-1
40+
\N val-null
41+
-- Rows read:
42+
2
43+
100
44+
4
45+
102
46+
47+
-- KeeperMap: set safe to cast
48+
0 val-0
49+
1 val-1
50+
-- KeeperMap: set not safe to cast
51+
0 val-0
52+
1 val-1
53+
-- KeeperMap: set with overflow
54+
0 val-0
55+
-- KeeperMap: set with cast failure
56+
-- KeeperMap: subquery safe to cast
57+
0 val-0
58+
1 val-1
59+
-- KeeperMap: subquery not safe to cast
60+
0 val-0
61+
1 val-1
62+
-- KeeperMap: subquery with overflow
63+
0 val-0
64+
-- KeeperMap: subquery with cast failure
65+
0 val-0
66+
-- Rows read:
67+
2
68+
2
69+
1
70+
4
71+
4
72+
102
73+
102
74+
75+
-- KeeperMap null: set without null
76+
0 val-0
77+
1 val-1
78+
-- KeeperMap null: set with null
79+
1 val-1
80+
\N val-null
81+
-- KeeperMap null: subquery without null
82+
0 val-0
83+
1 val-1
84+
-- KeeperMap null: subquery with null
85+
1 val-1
86+
\N val-null
87+
-- Rows read:
88+
2
89+
100
90+
4
91+
102

0 commit comments

Comments
 (0)