#include "TrieDictionary.h" #include #include #include #include #include #include #include #include #include #include #include #include #include #include "DictionaryBlockInputStream.h" #include "DictionaryFactory.h" namespace DB { namespace ErrorCodes { extern const int LOGICAL_ERROR; extern const int TYPE_MISMATCH; extern const int BAD_ARGUMENTS; extern const int DICTIONARY_IS_EMPTY; } static void validateKeyTypes(const DataTypes & key_types) { if (key_types.size() != 1) throw Exception{"Expected a single IP address", ErrorCodes::TYPE_MISMATCH}; const auto & actual_type = key_types[0]->getName(); if (actual_type != "UInt32" && actual_type != "FixedString(16)") throw Exception{"Key does not match, expected either UInt32 or FixedString(16)", ErrorCodes::TYPE_MISMATCH}; } /// Create IPAddress from 16 byte array converting to ipv4 if possible static Poco::Net::IPAddress ip4or6fromBytes(const uint8_t * data) { Poco::Net::IPAddress ipaddr(reinterpret_cast(data), IPV6_BINARY_LENGTH); // try to consider as ipv4 bool is_v4 = false; if (auto addr_v4 = IPv4ToBinary(ipaddr, is_v4); is_v4) return Poco::Net::IPAddress(reinterpret_cast(&addr_v4), IPV4_BINARY_LENGTH); return ipaddr; } TrieDictionary::TrieDictionary( const StorageID & dict_id_, const DictionaryStructure & dict_struct_, DictionarySourcePtr source_ptr_, const DictionaryLifetime dict_lifetime_, bool require_nonempty_) : IDictionaryBase(dict_id_) , dict_struct(dict_struct_) , source_ptr{std::move(source_ptr_)} , dict_lifetime(dict_lifetime_) , require_nonempty(require_nonempty_) , total_ip_length(0) , logger(&Poco::Logger::get("TrieDictionary")) { createAttributes(); loadData(); calculateBytesAllocated(); } TrieDictionary::~TrieDictionary() { } #define DECLARE(TYPE) \ void TrieDictionary::get##TYPE( \ const std::string & attribute_name, const Columns & key_columns, const DataTypes & key_types, ResultArrayType & out) const \ { \ validateKeyTypes(key_types); \ \ const auto & attribute = getAttribute(attribute_name); \ checkAttributeType(this, attribute_name, attribute.type, AttributeUnderlyingType::ut##TYPE); \ \ const auto null_value = std::get(attribute.null_values); \ \ getItemsImpl( \ attribute, \ key_columns, \ [&](const size_t row, const auto value) { out[row] = value; }, \ [&](const size_t) { return null_value; }); \ } DECLARE(UInt8) DECLARE(UInt16) DECLARE(UInt32) DECLARE(UInt64) DECLARE(UInt128) DECLARE(Int8) DECLARE(Int16) DECLARE(Int32) DECLARE(Int64) DECLARE(Float32) DECLARE(Float64) DECLARE(Decimal32) DECLARE(Decimal64) DECLARE(Decimal128) #undef DECLARE void TrieDictionary::getString( const std::string & attribute_name, const Columns & key_columns, const DataTypes & key_types, ColumnString * out) const { validateKeyTypes(key_types); const auto & attribute = getAttribute(attribute_name); checkAttributeType(this, attribute_name, attribute.type, AttributeUnderlyingType::utString); const auto & null_value = StringRef{std::get(attribute.null_values)}; getItemsImpl( attribute, key_columns, [&](const size_t, const StringRef value) { out->insertData(value.data, value.size); }, [&](const size_t) { return null_value; }); } #define DECLARE(TYPE) \ void TrieDictionary::get##TYPE( \ const std::string & attribute_name, \ const Columns & key_columns, \ const DataTypes & key_types, \ const PaddedPODArray & def, \ ResultArrayType & out) const \ { \ validateKeyTypes(key_types); \ \ const auto & attribute = getAttribute(attribute_name); \ checkAttributeType(this, attribute_name, attribute.type, AttributeUnderlyingType::ut##TYPE); \ \ getItemsImpl( \ attribute, \ key_columns, \ [&](const size_t row, const auto value) { out[row] = value; }, \ [&](const size_t row) { return def[row]; }); \ } DECLARE(UInt8) DECLARE(UInt16) DECLARE(UInt32) DECLARE(UInt64) DECLARE(UInt128) DECLARE(Int8) DECLARE(Int16) DECLARE(Int32) DECLARE(Int64) DECLARE(Float32) DECLARE(Float64) DECLARE(Decimal32) DECLARE(Decimal64) DECLARE(Decimal128) #undef DECLARE void TrieDictionary::getString( const std::string & attribute_name, const Columns & key_columns, const DataTypes & key_types, const ColumnString * const def, ColumnString * const out) const { validateKeyTypes(key_types); const auto & attribute = getAttribute(attribute_name); checkAttributeType(this, attribute_name, attribute.type, AttributeUnderlyingType::utString); getItemsImpl( attribute, key_columns, [&](const size_t, const StringRef value) { out->insertData(value.data, value.size); }, [&](const size_t row) { return def->getDataAt(row); }); } #define DECLARE(TYPE) \ void TrieDictionary::get##TYPE( \ const std::string & attribute_name, \ const Columns & key_columns, \ const DataTypes & key_types, \ const TYPE def, \ ResultArrayType & out) const \ { \ validateKeyTypes(key_types); \ \ const auto & attribute = getAttribute(attribute_name); \ checkAttributeType(this, attribute_name, attribute.type, AttributeUnderlyingType::ut##TYPE); \ \ getItemsImpl( \ attribute, key_columns, [&](const size_t row, const auto value) { out[row] = value; }, [&](const size_t) { return def; }); \ } DECLARE(UInt8) DECLARE(UInt16) DECLARE(UInt32) DECLARE(UInt64) DECLARE(UInt128) DECLARE(Int8) DECLARE(Int16) DECLARE(Int32) DECLARE(Int64) DECLARE(Float32) DECLARE(Float64) DECLARE(Decimal32) DECLARE(Decimal64) DECLARE(Decimal128) #undef DECLARE void TrieDictionary::getString( const std::string & attribute_name, const Columns & key_columns, const DataTypes & key_types, const String & def, ColumnString * const out) const { validateKeyTypes(key_types); const auto & attribute = getAttribute(attribute_name); checkAttributeType(this, attribute_name, attribute.type, AttributeUnderlyingType::utString); getItemsImpl( attribute, key_columns, [&](const size_t, const StringRef value) { out->insertData(value.data, value.size); }, [&](const size_t) { return StringRef{def}; }); } void TrieDictionary::has(const Columns & key_columns, const DataTypes & key_types, PaddedPODArray & out) const { validateKeyTypes(key_types); const auto & attribute = attributes.front(); switch (attribute.type) { case AttributeUnderlyingType::utUInt8: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utUInt16: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utUInt32: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utUInt64: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utUInt128: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utInt8: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utInt16: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utInt32: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utInt64: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utFloat32: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utFloat64: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utString: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utDecimal32: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utDecimal64: has(attribute, key_columns, out); break; case AttributeUnderlyingType::utDecimal128: has(attribute, key_columns, out); break; } } void TrieDictionary::createAttributes() { const auto size = dict_struct.attributes.size(); attributes.reserve(size); for (const auto & attribute : dict_struct.attributes) { attribute_index_by_name.emplace(attribute.name, attributes.size()); attributes.push_back(createAttributeWithType(attribute.underlying_type, attribute.null_value)); if (attribute.hierarchical) throw Exception{full_name + ": hierarchical attributes not supported for dictionary of type " + getTypeName(), ErrorCodes::TYPE_MISMATCH}; } } void TrieDictionary::loadData() { auto stream = source_ptr->loadAll(); stream->readPrefix(); /// created upfront to avoid excess allocations const auto keys_size = dict_struct.key->size(); ip_records.reserve(keys_size); const auto attributes_size = attributes.size(); while (const auto block = stream->read()) { const auto rows = block.rows(); element_count += rows; const auto key_column_ptrs = ext::map( ext::range(0, keys_size), [&](const size_t attribute_idx) { return block.safeGetByPosition(attribute_idx).column; }); const auto attribute_column_ptrs = ext::map(ext::range(0, attributes_size), [&](const size_t attribute_idx) { return block.safeGetByPosition(keys_size + attribute_idx).column; }); for (const auto row_idx : ext::range(0, rows)) { /// calculate key once per row const auto key_column = key_column_ptrs.front(); for (const auto attribute_idx : ext::range(0, attributes_size)) { const auto & attribute_column = *attribute_column_ptrs[attribute_idx]; auto & attribute = attributes[attribute_idx]; setAttributeValue(attribute, attribute_column[row_idx]); } size_t row_number = ip_records.size(); std::string addr_str(key_column->getDataAt(row_idx).toString()); size_t pos = addr_str.find('/'); if (pos != std::string::npos) { IPAddress addr(addr_str.substr(0, pos)); UInt8 prefix = std::stoi(addr_str.substr(pos + 1), nullptr, 10); addr = addr & IPAddress(prefix, addr.family()); ip_records.emplace_back(IPRecord{addr, prefix, row_number}); } else { IPAddress addr(addr_str); UInt8 prefix = addr.length() * 8; ip_records.emplace_back(IPRecord{addr, prefix, row_number}); } total_ip_length += ip_records.back().addr.length(); } } LOG_TRACE(logger, "{} ip records are read", ip_records.size()); std::sort(ip_records.begin(), ip_records.end(), [](const auto & a, const auto & b) { if (a.addr.family() != b.addr.family()) return a.addr.family() < b.addr.family(); // prefer IPs with more narrow subnet if (a.addr == b.addr) return a.prefix < b.prefix; return a.addr < b.addr; }); stream->readSuffix(); if (require_nonempty && 0 == element_count) throw Exception{full_name + ": dictionary source is empty and 'require_nonempty' property is set.", ErrorCodes::DICTIONARY_IS_EMPTY}; } template void TrieDictionary::addAttributeSize(const Attribute & attribute) { const auto & vec = std::get>(attribute.maps); bytes_allocated += sizeof(ContainerType) + (vec.capacity() * sizeof(T)); bucket_count = vec.size(); } void TrieDictionary::calculateBytesAllocated() { bytes_allocated += ip_records.size() * sizeof(ip_records.front()); bytes_allocated += total_ip_length; bytes_allocated += attributes.size() * sizeof(attributes.front()); for (const auto & attribute : attributes) { switch (attribute.type) { case AttributeUnderlyingType::utUInt8: addAttributeSize(attribute); break; case AttributeUnderlyingType::utUInt16: addAttributeSize(attribute); break; case AttributeUnderlyingType::utUInt32: addAttributeSize(attribute); break; case AttributeUnderlyingType::utUInt64: addAttributeSize(attribute); break; case AttributeUnderlyingType::utUInt128: addAttributeSize(attribute); break; case AttributeUnderlyingType::utInt8: addAttributeSize(attribute); break; case AttributeUnderlyingType::utInt16: addAttributeSize(attribute); break; case AttributeUnderlyingType::utInt32: addAttributeSize(attribute); break; case AttributeUnderlyingType::utInt64: addAttributeSize(attribute); break; case AttributeUnderlyingType::utFloat32: addAttributeSize(attribute); break; case AttributeUnderlyingType::utFloat64: addAttributeSize(attribute); break; case AttributeUnderlyingType::utDecimal32: addAttributeSize(attribute); break; case AttributeUnderlyingType::utDecimal64: addAttributeSize(attribute); break; case AttributeUnderlyingType::utDecimal128: addAttributeSize(attribute); break; case AttributeUnderlyingType::utString: { addAttributeSize(attribute); bytes_allocated += sizeof(Arena) + attribute.string_arena->size(); break; } } } } template void TrieDictionary::createAttributeImpl(Attribute & attribute, const Field & null_value) { attribute.null_values = T(null_value.get>()); attribute.maps.emplace>(); } TrieDictionary::Attribute TrieDictionary::createAttributeWithType(const AttributeUnderlyingType type, const Field & null_value) { Attribute attr{type, {}, {}, {}}; switch (type) { case AttributeUnderlyingType::utUInt8: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utUInt16: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utUInt32: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utUInt64: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utUInt128: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utInt8: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utInt16: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utInt32: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utInt64: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utFloat32: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utFloat64: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utDecimal32: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utDecimal64: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utDecimal128: createAttributeImpl(attr, null_value); break; case AttributeUnderlyingType::utString: { attr.null_values = null_value.get(); attr.maps.emplace>(); attr.string_arena = std::make_unique(); break; } } return attr; } template void TrieDictionary::getItemsImpl( const Attribute & attribute, const Columns & key_columns, ValueSetter && set_value, DefaultGetter && get_default) const { auto & vec = std::get>(attribute.maps); const auto first_column = key_columns.front(); const auto rows = first_column->size(); if (first_column->isNumeric()) { for (const auto i : ext::range(0, rows)) { auto addr = Poco::ByteOrder::toNetwork(UInt32(first_column->get64(i))); auto ipaddr = IPAddress(reinterpret_cast(&addr), IPV4_BINARY_LENGTH); auto found = lookupIPRecord(ipaddr); set_value(i, (found != ipRecordNotFound()) ? static_cast(vec[found->row]) : get_default(i)); } } else { for (const auto i : ext::range(0, rows)) { auto addr = first_column->getDataAt(i); if (addr.size != IPV6_BINARY_LENGTH) throw Exception("Expected key to be FixedString(16)", ErrorCodes::LOGICAL_ERROR); auto ipaddr = ip4or6fromBytes(reinterpret_cast(addr.data)); auto found = lookupIPRecord(ipaddr); set_value(i, (found != ipRecordNotFound()) ? static_cast(vec[found->row]) : get_default(i)); } } query_count.fetch_add(rows, std::memory_order_relaxed); } template void TrieDictionary::setAttributeValueImpl(Attribute & attribute, const T value) { auto & vec = std::get>(attribute.maps); vec.push_back(value); } void TrieDictionary::setAttributeValue(Attribute & attribute, const Field & value) { switch (attribute.type) { case AttributeUnderlyingType::utUInt8: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utUInt16: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utUInt32: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utUInt64: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utUInt128: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utInt8: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utInt16: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utInt32: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utInt64: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utFloat32: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utFloat64: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utDecimal32: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utDecimal64: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utDecimal128: return setAttributeValueImpl(attribute, value.get()); case AttributeUnderlyingType::utString: { const auto & string = value.get(); const auto * string_in_arena = attribute.string_arena->insert(string.data(), string.size()); return setAttributeValueImpl(attribute, StringRef{string_in_arena, string.size()}); } } } const TrieDictionary::Attribute & TrieDictionary::getAttribute(const std::string & attribute_name) const { const auto it = attribute_index_by_name.find(attribute_name); if (it == std::end(attribute_index_by_name)) throw Exception{full_name + ": no such attribute '" + attribute_name + "'", ErrorCodes::BAD_ARGUMENTS}; return attributes[it->second]; } template void TrieDictionary::has(const Attribute &, const Columns & key_columns, PaddedPODArray & out) const { const auto first_column = key_columns.front(); const auto rows = first_column->size(); if (first_column->isNumeric()) { for (const auto i : ext::range(0, rows)) { auto addr = Int32(first_column->get64(i)); auto ipaddr = IPAddress(reinterpret_cast(&addr), IPV4_BINARY_LENGTH); auto found = lookupIPRecord(ipaddr); out[i] = (found != ipRecordNotFound()); } } else { for (const auto i : ext::range(0, rows)) { auto addr = first_column->getDataAt(i); if (unlikely(addr.size != 16)) throw Exception("Expected key to be FixedString(16)", ErrorCodes::LOGICAL_ERROR); auto ipaddr = ip4or6fromBytes(reinterpret_cast(addr.data)); auto found = lookupIPRecord(ipaddr); out[i] = (found != ipRecordNotFound()); } } query_count.fetch_add(rows, std::memory_order_relaxed); } Columns TrieDictionary::getKeyColumns() const { auto ip_column = ColumnFixedString::create(IPV6_BINARY_LENGTH); auto mask_column = ColumnVector::create(); for (const auto & record : ip_records) { auto ip_array = IPv6ToBinary(record.addr); ip_column->insertData(ip_array.data(), IPV6_BINARY_LENGTH); mask_column->insertValue(record.prefix); } return {std::move(ip_column), std::move(mask_column)}; } BlockInputStreamPtr TrieDictionary::getBlockInputStream(const Names & column_names, size_t max_block_size) const { using BlockInputStreamType = DictionaryBlockInputStream; auto get_keys = [](const Columns & columns, const std::vector & dict_attributes) { const auto & attr = dict_attributes.front(); return ColumnsWithTypeAndName( {ColumnWithTypeAndName(columns.front(), std::make_shared(IPV6_BINARY_LENGTH), attr.name)}); }; auto get_view = [](const Columns & columns, const std::vector & dict_attributes) { auto column = ColumnString::create(); const auto & ip_column = assert_cast(*columns.front()); const auto & mask_column = assert_cast &>(*columns.back()); char buffer[48]; for (size_t row : ext::range(0, ip_column.size())) { UInt8 mask = mask_column.getElement(row); char * ptr = buffer; formatIPv6(reinterpret_cast(ip_column.getDataAt(row).data), ptr); *(ptr - 1) = '/'; ptr = itoa(mask, ptr); column->insertData(buffer, ptr - buffer); } return ColumnsWithTypeAndName{ ColumnWithTypeAndName(std::move(column), std::make_shared(), dict_attributes.front().name)}; }; return std::make_shared( shared_from_this(), max_block_size, getKeyColumns(), column_names, std::move(get_keys), std::move(get_view)); } int TrieDictionary::matchIPAddrWithRecord(const IPAddress & ipaddr, const IPRecord & record) const { if (ipaddr.family() != record.addr.family()) return ipaddr.family() < record.addr.family() ? -1 : 1; auto masked_ipaddr = ipaddr & IPAddress(record.prefix, record.addr.family()); if (masked_ipaddr < record.addr) return -1; if (masked_ipaddr == record.addr) return 0; return 1; } TrieDictionary::IPRecordConstIt TrieDictionary::ipRecordNotFound() const { return ip_records.end(); } TrieDictionary::IPRecordConstIt TrieDictionary::lookupIPRecord(const IPAddress & target) const { if (ip_records.empty()) return ipRecordNotFound(); auto comp = [&](const IPAddress & needle, const IPRecord & record) -> bool { return matchIPAddrWithRecord(needle, record) < 0; }; auto next_it = std::upper_bound(ip_records.begin(), ip_records.end(), target, comp); if (next_it == ip_records.begin()) return ipRecordNotFound(); auto found = next_it - 1; if (matchIPAddrWithRecord(target, *found) == 0) return found; return ipRecordNotFound(); } void registerDictionaryTrie(DictionaryFactory & factory) { auto create_layout = [=](const std::string &, const DictionaryStructure & dict_struct, const Poco::Util::AbstractConfiguration & config, const std::string & config_prefix, DictionarySourcePtr source_ptr) -> DictionaryPtr { if (!dict_struct.key) throw Exception{"'key' is required for dictionary of layout 'ip_trie'", ErrorCodes::BAD_ARGUMENTS}; const auto dict_id = StorageID::fromDictionaryConfig(config, config_prefix); const DictionaryLifetime dict_lifetime{config, config_prefix + ".lifetime"}; const bool require_nonempty = config.getBool(config_prefix + ".require_nonempty", false); // This is specialised trie for storing IPv4 and IPv6 prefixes. return std::make_unique(dict_id, dict_struct, std::move(source_ptr), dict_lifetime, require_nonempty); }; factory.registerLayout("ip_trie", create_layout, true); } }