Skip to content

Commit

Permalink
IPAddressDictionary fixed, test fixed, porting ClickHouse/ClickHouse#…
Browse files Browse the repository at this point in the history 
…43221
  • Loading branch information
lijianan committed Dec 6, 2023
1 parent 3f218c8 commit 5b51ee1
Show file tree
Hide file tree
Showing 2 changed files with 123 additions and 12 deletions.
31 changes: 20 additions & 11 deletions src/Dictionaries/IPAddressDictionary.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#include <DataTypes/DataTypeIPv4andIPv6.h>
#include <Poco/ByteOrder.h>
#include <Common/formatIPv6.h>
#include "Core/Types.h"
#include <Core/Types.h>
#include <base/itoa.h>
#include <base/map.h>
#include <base/range.h>
Expand Down Expand Up @@ -131,12 +131,13 @@ static void validateKeyTypes(const DataTypes & key_types)
if (key_types.empty() || key_types.size() > 2)
throw Exception(ErrorCodes::TYPE_MISMATCH, "Expected a single IP address or IP with mask");

const auto * key_ipv4type = typeid_cast<const DataTypeIPv4 *>(key_types[0].get());
const auto * key_ipv6type = typeid_cast<const DataTypeIPv6 *>(key_types[0].get());
TypeIndex type_id = key_types[0]->getTypeId();
const auto * key_string = typeid_cast<const DataTypeFixedString *>(key_types[0].get());

if (key_ipv4type == nullptr && key_ipv6type == nullptr)
if (type_id != TypeIndex::IPv4 && type_id != TypeIndex::UInt32 && type_id != TypeIndex::IPv6 && !(key_string && key_string->getN() == IPV6_BINARY_LENGTH))
throw Exception(ErrorCodes::TYPE_MISMATCH,
"Key does not match, expected either `IPv4` or `IPv6`");
"Key does not match, expected either ipv4 (or uint32) or ipv6 (or fixed_string(16))");


if (key_types.size() > 1)
{
Expand Down Expand Up @@ -298,7 +299,9 @@ ColumnUInt8::Ptr IPAddressDictionary::hasKeys(const Columns & key_columns, const

size_t keys_found = 0;

if (first_column->getDataType() == TypeIndex::IPv4)
TypeIndex type_id = first_column->getDataType();

if (type_id == TypeIndex::IPv4 || type_id == TypeIndex::UInt32)
{
uint8_t addrv6_buf[IPV6_BINARY_LENGTH];
for (const auto i : collections::range(0, rows))
Expand All @@ -309,11 +312,13 @@ ColumnUInt8::Ptr IPAddressDictionary::hasKeys(const Columns & key_columns, const
keys_found += out[i];
}
}
else if (first_column->getDataType() == TypeIndex::IPv6)
else if (type_id == TypeIndex::IPv6 || type_id == TypeIndex::FixedString)
{
for (const auto i : collections::range(0, rows))
{
auto addr = first_column->getDataAt(i);
if (addr.size != IPV6_BINARY_LENGTH)
throw Exception(ErrorCodes::TYPE_MISMATCH, "Expected key fixed_string(16)");
auto found = tryLookupIPv6(reinterpret_cast<const uint8_t *>(addr.data));
out[i] = (found != ipNotFound());
keys_found += out[i];
Expand Down Expand Up @@ -538,7 +543,7 @@ template <>
void IPAddressDictionary::addAttributeSize<String>(const Attribute & attribute)
{
addAttributeSize<StringRef>(attribute);
bytes_allocated += sizeof(Arena) + attribute.string_arena->size();
bytes_allocated += sizeof(Arena) + attribute.string_arena->allocatedBytes();
}

void IPAddressDictionary::calculateBytesAllocated()
Expand Down Expand Up @@ -711,7 +716,9 @@ void IPAddressDictionary::getItemsImpl(

size_t keys_found = 0;

if (first_column->getDataType() == TypeIndex::IPv4)
TypeIndex type_id = first_column->getDataType();

if (type_id == TypeIndex::IPv4 || type_id == TypeIndex::UInt32)
{
uint8_t addrv6_buf[IPV6_BINARY_LENGTH];
for (const auto i : collections::range(0, rows))
Expand All @@ -728,11 +735,13 @@ void IPAddressDictionary::getItemsImpl(
set_value(i, default_value_extractor[i]);
}
}
else if (first_column->getDataType() == TypeIndex::IPv6)
else if (type_id == TypeIndex::IPv6 || type_id == TypeIndex::FixedString)
{
for (const auto i : collections::range(0, rows))
{
auto addr = first_column->getDataAt(i);
if (addr.size != IPV6_BINARY_LENGTH)
throw Exception(ErrorCodes::TYPE_MISMATCH, "Expected key to be fixed_string(16)");
auto found = tryLookupIPv6(reinterpret_cast<const uint8_t *>(addr.data));
if (found != ipNotFound())
{
Expand All @@ -744,7 +753,7 @@ void IPAddressDictionary::getItemsImpl(
}
}
else
throw Exception(ErrorCodes::TYPE_MISMATCH, "Expected key to be IPv4 or IPv6");
throw Exception(ErrorCodes::TYPE_MISMATCH, "Expected key to be ipv4 (or uint32) or ipv6 (or fixed_string(16))");

query_count.fetch_add(rows, std::memory_order_relaxed);
found_count.fetch_add(keys_found, std::memory_order_relaxed);
Expand Down
104 changes: 103 additions & 1 deletion tests/stream/test_stream_smoke/0002_ingest_and_ddl_case.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,108 @@
["dev9", "ca", 100, "2020-01-01 11:11:11"]
]}
]
}
},
{
"id": 11,
"tags": [
"ip_trie_dictionary"
],
"name": "test ip_trie dictionary data ingest and dict_get function",
"description": "test ip_trie dictionary data ingest and dict_get function",
"steps": [
{
"statements": [
{
"client": "python",
"query_type": "table",
"query": "drop stream if exists test2_stream"
},
{
"client": "python",
"query_type": "table",
"exist_wait": 2,
"wait": 1,
"query": "
create stream if not exists geoip_t (
cidr string,
latitude float64,
longitude float64,
country_code string,
state string,
city string
) engine = MergeTree"
},
{
"client": "python",
"query_type": "table",
"depends_on_stream" : "geoip_t",
"query": "insert into geoip_t values ('188.166.84.125',52.3759,4.8975,'NL','North Holland','Amsterdam')"
},
{
"client": "python",
"query_type": "table",
"query": "drop dictionary geoip;"
},
{
"client": "python",
"query_type": "table",
"query": "
CREATE DICTIONARY geoip (
cidr string,
latitude float64,
longitude float64,
country_code string,
state string,
city string
) PRIMARY KEY cidr
SOURCE (ClickHouse (table 'geoip_t'))
LIFETIME (MIN 300 MAX 360)
LAYOUT (IP_TRIE());"
},
{
"client": "python",
"query_type": "table",
"query_id": "211",
"query": "select dict_get('geoip', 'latitude', to_ipv4('188.166.84.125'))"
},
{
"client": "python",
"query_type": "table",
"query_id": "212",
"query": "select dict_get_or_default('geoip', 'latitude', to_ipv4('188.166.84.111'), 1.0)"
},
{
"client": "python",
"query_type": "table",
"query_id": "213",
"query": "select dict_get_or_default('geoip', 'latitude', '188.166.84.111', 1.0)"
}
]
}
],
"expected_results": [
{
"query_id": "211",
"expected_results": [
[
"52.3759"
]
]
},
{
"query_id": "212",
"expected_results": [
[
"1.0"
]
]
},
{
"query_id": "213",
"expected_results":"error_code:53"
}

]
}
]
}

0 comments on commit 5b51ee1

Please sign in to comment.