blob: 3f73e6736a0a3a52903ea5e74f8ba8e5c5fe2088 [file]
/*
* 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.
*/
#include "paimon/common/file_index/rangebitmap/range_bitmap.h"
#include <algorithm>
#include "fmt/format.h"
#include "paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary.h"
#include "paimon/common/file_index/rangebitmap/dictionary/key_factory.h"
#include "paimon/common/file_index/rangebitmap/utils/literal_serialization_utils.h"
#include "paimon/common/io/data_output_stream.h"
#include "paimon/common/io/memory_segment_output_stream.h"
#include "paimon/common/memory/memory_segment_utils.h"
#include "paimon/common/utils/field_type_utils.h"
#include "paimon/io/data_input_stream.h"
#include "paimon/memory/bytes.h"
namespace paimon {
Result<std::unique_ptr<RangeBitmap>> RangeBitmap::Create(
const std::shared_ptr<InputStream>& input_stream, const int64_t offset,
const FieldType field_type, const std::shared_ptr<MemoryPool>& pool) {
PAIMON_RETURN_NOT_OK(input_stream->Seek(offset, SeekOrigin::FS_SEEK_SET));
const auto data_in = std::make_shared<DataInputStream>(input_stream);
PAIMON_ASSIGN_OR_RAISE(int32_t header_length, data_in->ReadValue<int32_t>());
PAIMON_ASSIGN_OR_RAISE(int8_t version, data_in->ReadValue<int8_t>());
if (version != kCurrentVersion) {
return Status::Invalid(fmt::format("RangeBitmap unsupported version: {}", version));
}
PAIMON_ASSIGN_OR_RAISE(int32_t rid, data_in->ReadValue<int32_t>());
PAIMON_ASSIGN_OR_RAISE(int32_t cardinality, data_in->ReadValue<int32_t>());
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<KeyFactory> shared_key_factory,
KeyFactory::Create(field_type));
PAIMON_ASSIGN_OR_RAISE(LiteralSerDeUtils::Deserializer key_deserializer,
LiteralSerDeUtils::CreateValueReader(field_type));
auto min = Literal{field_type};
auto max = Literal{field_type};
if (cardinality > 0) {
PAIMON_ASSIGN_OR_RAISE(min, key_deserializer(data_in, pool.get()));
PAIMON_ASSIGN_OR_RAISE(max, key_deserializer(data_in, pool.get()));
}
PAIMON_ASSIGN_OR_RAISE(int32_t dictionary_length, data_in->ReadValue<int32_t>());
auto dictionary_offset = static_cast<int32_t>(offset + sizeof(int32_t) + header_length);
int32_t bsi_offset = dictionary_offset + dictionary_length;
return std::unique_ptr<RangeBitmap>(new RangeBitmap(rid, cardinality, dictionary_offset,
bsi_offset, min, max, shared_key_factory,
input_stream, pool));
}
Status RangeBitmap::Not(RoaringBitmap32* out) {
out->Flip(0, rid_);
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 is_not_null, this->IsNotNull());
*out &= is_not_null;
return Status::OK();
}
Result<RoaringBitmap32> RangeBitmap::Eq(const Literal& key) {
if (cardinality_ <= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(int32_t min_compare, key.CompareTo(min_));
PAIMON_ASSIGN_OR_RAISE(int32_t max_compare, key.CompareTo(max_));
PAIMON_ASSIGN_OR_RAISE(auto* bit_slice_ptr, this->GetBitSliceIndex());
if (min_compare == 0 && max_compare == 0) {
return bit_slice_ptr->IsNotNull({});
}
if (min_compare < 0 || max_compare > 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(auto* dictionary, this->GetDictionary());
PAIMON_ASSIGN_OR_RAISE(int32_t code, dictionary->Find(key));
if (code < 0) {
return RoaringBitmap32();
}
return bit_slice_ptr->Eq(code);
}
Result<RoaringBitmap32> RangeBitmap::Neq(const Literal& key) {
if (cardinality_ <= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, Eq(key));
PAIMON_RETURN_NOT_OK(Not(&result));
return result;
}
Result<RoaringBitmap32> RangeBitmap::Lt(const Literal& key) {
if (cardinality_ <= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(int32_t min_compare, key.CompareTo(min_));
PAIMON_ASSIGN_OR_RAISE(int32_t max_compare, key.CompareTo(max_));
if (max_compare > 0) {
return IsNotNull();
}
if (min_compare <= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, Gte(key));
PAIMON_RETURN_NOT_OK(Not(&result));
return result;
}
Result<RoaringBitmap32> RangeBitmap::Lte(const Literal& key) {
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 lt_result, Lt(key));
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 eq_result, Eq(key));
lt_result |= eq_result;
return lt_result;
}
Result<RoaringBitmap32> RangeBitmap::Gt(const Literal& key) {
if (cardinality_ <= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(int32_t max_compare, key.CompareTo(max_));
if (max_compare >= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(int32_t min_compare, key.CompareTo(min_));
if (min_compare < 0) {
return IsNotNull();
}
PAIMON_ASSIGN_OR_RAISE(auto* dictionary, this->GetDictionary());
PAIMON_ASSIGN_OR_RAISE(int32_t code, dictionary->Find(key));
PAIMON_ASSIGN_OR_RAISE(auto* bit_slice_ptr, this->GetBitSliceIndex());
if (code >= 0) {
return bit_slice_ptr->Gt(code);
}
return bit_slice_ptr->Gte(-code - 1);
}
Result<RoaringBitmap32> RangeBitmap::Gte(const Literal& key) {
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 gt_result, Gt(key));
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 eq_result, Eq(key));
gt_result |= eq_result;
return gt_result;
}
Result<RoaringBitmap32> RangeBitmap::In(const std::vector<Literal>& keys) {
if (cardinality_ <= 0) {
return RoaringBitmap32();
}
RoaringBitmap32 result{};
for (const auto& key : keys) {
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 bitmap, Eq(key));
result |= bitmap;
}
return result;
}
Result<RoaringBitmap32> RangeBitmap::NotIn(const std::vector<Literal>& keys) {
if (cardinality_ <= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, In(keys));
PAIMON_RETURN_NOT_OK(Not(&result));
return result;
}
Result<RoaringBitmap32> RangeBitmap::IsNull() {
if (cardinality_ <= 0) {
if (rid_ > 0) {
RoaringBitmap32 result;
result.AddRange(0, rid_);
return result;
}
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 non_null_bitmap, IsNotNull());
non_null_bitmap.Flip(0, rid_);
return non_null_bitmap;
}
Result<RoaringBitmap32> RangeBitmap::IsNotNull() {
if (cardinality_ <= 0) {
return RoaringBitmap32();
}
PAIMON_ASSIGN_OR_RAISE(auto* bit_slice_ptr, this->GetBitSliceIndex());
PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 result, bit_slice_ptr->IsNotNull({}));
return result;
}
RangeBitmap::RangeBitmap(const int32_t rid, const int32_t cardinality,
const int32_t dictionary_offset, const int32_t bsi_offset,
const Literal& min, const Literal& max,
const std::shared_ptr<KeyFactory>& key_factory,
const std::shared_ptr<InputStream>& input_stream,
const std::shared_ptr<MemoryPool>& pool)
: pool_(pool),
rid_(rid),
cardinality_(cardinality),
bsi_offset_(bsi_offset),
dictionary_offset_(dictionary_offset),
min_(min),
max_(max),
key_factory_(key_factory),
input_stream_(input_stream),
bsi_(nullptr),
dictionary_(nullptr) {}
Result<BitSliceIndexBitmap*> RangeBitmap::GetBitSliceIndex() {
if (bsi_ == nullptr) {
PAIMON_ASSIGN_OR_RAISE(bsi_,
BitSliceIndexBitmap::Create(input_stream_, bsi_offset_, pool_));
}
return bsi_.get();
}
Result<Dictionary*> RangeBitmap::GetDictionary() {
if (dictionary_ == nullptr) {
PAIMON_ASSIGN_OR_RAISE(
dictionary_, ChunkedDictionary::Create(key_factory_->GetFieldType(), input_stream_,
dictionary_offset_, pool_));
}
return dictionary_.get();
}
Result<std::unique_ptr<RangeBitmap::Appender>> RangeBitmap::Appender::Create(
const std::shared_ptr<KeyFactory>& factory, int64_t limited_serialized_size_in_bytes,
const std::shared_ptr<MemoryPool>& pool) {
return std::unique_ptr<Appender>(new Appender(factory, limited_serialized_size_in_bytes, pool));
}
RangeBitmap::Appender::Appender(const std::shared_ptr<KeyFactory>& factory,
const int64_t limited_serialized_size_in_bytes,
const std::shared_ptr<MemoryPool>& pool)
: pool_(pool),
rid_(0),
bitmaps_(LiteralComparator(factory)),
factory_(factory),
chunk_size_bytes_limit_(limited_serialized_size_in_bytes) {}
void RangeBitmap::Appender::Append(const Literal& key) {
if (!key.IsNull()) {
bitmaps_[key].Add(rid_);
}
rid_++;
}
Result<PAIMON_UNIQUE_PTR<Bytes>> RangeBitmap::Appender::Serialize() const {
int32_t code = 0;
const int32_t max_code = bitmaps_.empty() ? 0 : static_cast<int32_t>(bitmaps_.size() - 1);
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BitSliceIndexBitmap::Appender> bsi_appender,
BitSliceIndexBitmap::Appender::Create(0, max_code, pool_));
if (chunk_size_bytes_limit_ >= std::numeric_limits<int32_t>::max()) {
return Status::Invalid(fmt::format(
"Chunk size cannot be larger than 2GB, current bytes: {}", chunk_size_bytes_limit_));
}
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ChunkedDictionary::Appender> dictionary_appender,
ChunkedDictionary::Appender::Create(
factory_, static_cast<int32_t>(chunk_size_bytes_limit_), pool_));
for (const auto& [key, bitmap] : bitmaps_) {
PAIMON_RETURN_NOT_OK(dictionary_appender->AppendSorted(key, code));
for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) {
PAIMON_RETURN_NOT_OK(bsi_appender->Append(*it, code));
}
code++;
}
PAIMON_ASSIGN_OR_RAISE(LiteralSerDeUtils::Serializer serializer,
LiteralSerDeUtils::CreateValueWriter(factory_->GetFieldType()));
Literal min{factory_->GetFieldType()};
Literal max{factory_->GetFieldType()};
if (!bitmaps_.empty()) {
min = bitmaps_.begin()->first;
max = bitmaps_.rbegin()->first;
}
PAIMON_ASSIGN_OR_RAISE(int32_t min_size, LiteralSerDeUtils::GetSerializedSizeInBytes(min));
PAIMON_ASSIGN_OR_RAISE(int32_t max_size, LiteralSerDeUtils::GetSerializedSizeInBytes(max));
int32_t header_size = 0;
header_size += sizeof(int8_t); // version
header_size += sizeof(int32_t); // rid
header_size += sizeof(int32_t); // cardinality
header_size += min.IsNull() ? 0 : min_size; // min literal size
header_size += max.IsNull() ? 0 : max_size; // max literal size
header_size += sizeof(int32_t); // dictionary length
PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> dictionary_bytes,
dictionary_appender->Serialize());
auto dictionary_length = static_cast<int32_t>(dictionary_bytes->size());
PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> bsi_bytes, bsi_appender->Serialize());
size_t bsi_length = bsi_bytes->size();
const auto data_output_stream = std::make_shared<MemorySegmentOutputStream>(
MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_);
data_output_stream->WriteValue<int32_t>(header_size);
data_output_stream->WriteValue<int8_t>(kCurrentVersion);
data_output_stream->WriteValue<int32_t>(rid_);
data_output_stream->WriteValue<int32_t>(static_cast<int32_t>(bitmaps_.size()));
if (!min.IsNull()) {
PAIMON_RETURN_NOT_OK(serializer(data_output_stream, min));
}
if (!max.IsNull()) {
PAIMON_RETURN_NOT_OK(serializer(data_output_stream, max));
}
data_output_stream->WriteValue<int32_t>(dictionary_length);
data_output_stream->Write(dictionary_bytes->data(), dictionary_length);
data_output_stream->Write(bsi_bytes->data(), bsi_length);
return MemorySegmentUtils::CopyToBytes(data_output_stream->Segments(), 0,
static_cast<int32_t>(data_output_stream->CurrentSize()),
pool_.get());
}
} // namespace paimon