Skip to content

feat(file-index): support additional data types #271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zjw1111 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(file-index): support additional data types #271

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/paimon/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -572,6 +572,7 @@ if(PAIMON_BUILD_TESTS)
common/file_index/rangebitmap/dictionary/chunked_dictionary_test.cpp
common/file_index/rangebitmap/range_bitmap_file_index_test.cpp
common/file_index/rangebitmap/range_bitmap_io_test.cpp
common/file_index/rangebitmap/range_bitmap_type_adapter_test.cpp
common/file_index/bloomfilter/bloom_filter_file_index_test.cpp
common/file_index/bloomfilter/fast_hash_test.cpp
common/global_index/complete_index_score_batch_reader_test.cpp
Expand Down
4 changes: 3 additions & 1 deletion src/paimon/common/file_index/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,13 @@ set(PAIMON_FILE_INDEX_SRC
rangebitmap/dictionary/chunked_dictionary.cpp
rangebitmap/dictionary/fixed_length_chunk.cpp
rangebitmap/dictionary/key_factory.cpp
rangebitmap/dictionary/variable_length_chunk.cpp
rangebitmap/utils/literal_serialization_utils.cpp
rangebitmap/bit_slice_index_bitmap.cpp
rangebitmap/range_bitmap.cpp
rangebitmap/range_bitmap_file_index.cpp
rangebitmap/range_bitmap_file_index_factory.cpp)
rangebitmap/range_bitmap_file_index_factory.cpp
rangebitmap/range_bitmap_type_adapter.cpp)

add_paimon_lib(paimon_file_index
SOURCES
Expand Down
32 changes: 32 additions & 0 deletions src/paimon/common/file_index/bitmap/bitmap_file_index_meta.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,13 @@

#include "paimon/common/file_index/bitmap/bitmap_file_index_meta.h"

#include <cstdint>
#include <string>
#include <utility>

#include "fmt/format.h"
#include "paimon/common/utils/field_type_utils.h"
#include "paimon/common/utils/math.h"
#include "paimon/defs.h"
#include "paimon/io/data_input_stream.h"
#include "paimon/memory/bytes.h"
Expand Down Expand Up @@ -80,6 +82,18 @@ Result<std::function<void(const Literal&)>> BitmapFileIndexMeta::GetValueWriter(
[output_stream](const Literal& literal) -> void {
output_stream->WriteValue<int64_t>(literal.GetValue<int64_t>());
});
case FieldType::FLOAT:
return std::function<void(const Literal&)>(
[output_stream](const Literal& literal) -> void {
const float value = CanonicalizeFloatingPoint(literal.GetValue<float>());
output_stream->WriteValue<float>(value);
});
case FieldType::DOUBLE:
return std::function<void(const Literal&)>(
[output_stream](const Literal& literal) -> void {
const double value = CanonicalizeFloatingPoint(literal.GetValue<double>());
output_stream->WriteValue<double>(value);
});
case FieldType::STRING:
return std::function<void(const Literal&)>(
[output_stream](const Literal& literal) -> void {
Expand Down Expand Up @@ -155,6 +169,24 @@ Result<std::function<Result<Literal>()>> BitmapFileIndexMeta::GetValueReader(
};
return func;
}
case FieldType::FLOAT: {
std::function<Result<Literal>()> func = [&in, move_body_start,
this]() -> Result<Literal> {
PAIMON_ASSIGN_OR_RAISE(float value,
ReadAndMoveBodyStart<float>(in, move_body_start));
return Literal(value);
};
return func;
}
case FieldType::DOUBLE: {
std::function<Result<Literal>()> func = [&in, move_body_start,
this]() -> Result<Literal> {
PAIMON_ASSIGN_OR_RAISE(double value,
ReadAndMoveBodyStart<double>(in, move_body_start));
return Literal(value);
};
return func;
}
case FieldType::DATE: {
std::function<Result<Literal>()> func = [&in, move_body_start,
this]() -> Result<Literal> {
Expand Down
248 changes: 248 additions & 0 deletions src/paimon/common/file_index/bitmap/bitmap_file_index_test.cpp
View file
Open in desktop

Large diffs are not rendered by default.

21 changes: 19 additions & 2 deletions src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
#include "paimon/common/utils/checked_cast.h"
#include "paimon/common/utils/date_time_utils.h"
#include "paimon/common/utils/field_type_utils.h"
#include "paimon/data/decimal.h"
#include "paimon/data/timestamp.h"
#include "paimon/defs.h"
#include "paimon/file_index/bitmap_index_result.h"
Expand Down Expand Up @@ -252,10 +253,26 @@ Result<BitSliceIndexBitmapFileIndex::ValueMapperType> BitSliceIndexBitmapFileInd
return literal.GetValue<Timestamp>().ToMicrosecond();
});
}
case FieldType::DECIMAL:
return BitSliceIndexBitmapFileIndex::ValueMapperType(
[](const Literal& literal) -> Result<int64_t> {
if (literal.IsNull()) {
return Status::Invalid(
"literal cannot be null when GetValue in BitSliceIndexBitmapFileIndex");
}
const auto value = literal.GetValue<Decimal>();
if (value.Value() < std::numeric_limits<int64_t>::min() ||
value.Value() > std::numeric_limits<int64_t>::max()) {
return Status::Invalid(fmt::format(
"decimal unscaled value {} does not fit in int64 for bsi index",
value.ToString()));
}
return value.ToUnscaledLong();
});
default:
// TODO(xinyu.lxy): support decimal
return Status::Invalid(
"BitSliceIndexBitmapFileIndex only support TINYINT/SMALLINT/INT/BIGINT/DATE");
"BitSliceIndexBitmapFileIndex only support "
"TINYINT/SMALLINT/INT/BIGINT/DATE/TIMESTAMP/DECIMAL");
}
}

Expand Down
22 changes: 22 additions & 0 deletions src/paimon/common/file_index/bsi/bit_slice_index_bitmap_file_index_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
#include "gtest/gtest.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/field_type_utils.h"
#include "paimon/data/decimal.h"
#include "paimon/data/timestamp.h"
#include "paimon/defs.h"
#include "paimon/file_index/bitmap_index_result.h"
Expand Down Expand Up @@ -420,6 +421,27 @@ TEST_F(BitSliceIndexBitmapIndexReaderTest, TestTimestampType) {
"literal cannot be null when GetValue in BitSliceIndexBitmapFileIndex");
}

TEST_F(BitSliceIndexBitmapIndexReaderTest, TestDecimalType) {
const auto type = arrow::decimal128(10, 2);
ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR<Bytes> index_bytes,
WriteIndex(type, R"([["1.00"], ["2.50"], [null], ["-1.25"], ["2.50"]])"));
auto input_stream =
std::make_shared<ByteArrayInputStream>(index_bytes->data(), index_bytes->size());
BitSliceIndexBitmapFileIndex file_index({});
ASSERT_OK_AND_ASSIGN(
auto reader, file_index.CreateReader(CreateArrowSchema(type).get(), 0, index_bytes->size(),
input_stream, pool_));

CheckResult(reader->VisitEqual(Literal(Decimal(10, 2, 250))).value(), {1, 4});
CheckResult(reader->VisitGreaterThan(Literal(Decimal(10, 2, 100))).value(), {1, 4});
CheckResult(reader->VisitLessThan(Literal(Decimal(10, 2, 0))).value(), {3});
CheckResult(reader->VisitIsNull().value(), {2});

// test invalid case for decimal128(20, 0) which exceeds int64 range
ASSERT_NOK_WITH_MSG(WriteIndex(arrow::decimal128(20, 0), R"([["9223372036854775808"]])"),
"does not fit in int64 for bsi index");
}

TEST_F(BitSliceIndexBitmapIndexReaderTest, TestUnInvalidType) {
std::vector<char> index_bytes = {
1, 0, 0, 0, 5, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 58,
Expand Down
6 changes: 1 addition & 5 deletions src/paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,14 +89,11 @@ Result<std::shared_ptr<Chunk>> ChunkedDictionary::GetChunk(int32_t index) {
if (index < 0 || index >= size_) {
return Status::Invalid(fmt::format("Invalid chunk index: {}", index));
}
if (offsets_bytes_ == nullptr || chunks_bytes_ == nullptr) {
if (offsets_bytes_ == nullptr) {
PAIMON_RETURN_NOT_OK(input_stream_->Seek(body_offset_, FS_SEEK_SET));
auto offsets = Bytes::AllocateBytes(offsets_length_, pool_.get());
PAIMON_RETURN_NOT_OK(input_stream_->Read(offsets->data(), offsets_length_));
offsets_bytes_ = std::move(offsets);
auto chunks = Bytes::AllocateBytes(chunks_length_, pool_.get());
PAIMON_RETURN_NOT_OK(input_stream_->Read(chunks->data(), chunks_length_));
chunks_bytes_ = std::move(chunks);
}
if (chunks_cache_[index]) {
return chunks_cache_[index];
Expand Down Expand Up @@ -244,6 +241,5 @@ ChunkedDictionary::ChunkedDictionary(const std::shared_ptr<InputStream>& input_s
chunks_length_(chunks_length),
body_offset_(body_offset),
offsets_bytes_(nullptr),
chunks_bytes_(nullptr),
chunks_cache_(std::vector<std::shared_ptr<Chunk>>(size)) {}
} // namespace paimon
1 change: 0 additions & 1 deletion src/paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,6 @@ class ChunkedDictionary final : public Dictionary {

// for lazy loading
PAIMON_UNIQUE_PTR<Bytes> offsets_bytes_;
PAIMON_UNIQUE_PTR<Bytes> chunks_bytes_;

// mmap chunks cache
std::vector<std::shared_ptr<Chunk>> chunks_cache_;
Expand Down
22 changes: 19 additions & 3 deletions src/paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -456,9 +456,25 @@ TEST_F(ChunkedDictionaryTest, TestKeyFactoryUnsupportedType) {
"Unsupported field type for KeyFactory: BINARY");
}

TEST_F(ChunkedDictionaryTest, TestStringKeyFactoryNotImplemented) {
ASSERT_NOK_WITH_MSG(KeyFactory::Create(FieldType::STRING),
"Unsupported field type for KeyFactory: STRING");
TEST_F(ChunkedDictionaryTest, TestStringKeyFactory) {
ASSERT_OK_AND_ASSIGN(auto key_factory, KeyFactory::Create(FieldType::STRING));
ASSERT_OK_AND_ASSIGN(auto appender,
ChunkedDictionary::Appender::Create(key_factory, 12, pool_));
ASSERT_OK(appender->AppendSorted(Literal(FieldType::STRING, "apple", 5), 0));
ASSERT_OK(appender->AppendSorted(Literal(FieldType::STRING, "banana", 6), 1));
ASSERT_OK(appender->AppendSorted(Literal(FieldType::STRING, "pear", 4), 2));
ASSERT_OK_AND_ASSIGN(auto bytes, appender->Serialize());
auto input_stream = std::make_shared<ByteArrayInputStream>(bytes->data(), bytes->size());
ASSERT_OK_AND_ASSIGN(auto dict,
ChunkedDictionary::Create(FieldType::STRING, input_stream, 0, pool_));

ASSERT_OK_AND_ASSIGN(int32_t banana_code, dict->Find(Literal(FieldType::STRING, "banana", 6)));
ASSERT_EQ(banana_code, 1);
ASSERT_OK_AND_ASSIGN(Literal pear, dict->Find(2));
ASSERT_EQ(pear.GetValue<std::string>(), "pear");
ASSERT_OK_AND_ASSIGN(int32_t between_code,
dict->Find(Literal(FieldType::STRING, "blueberry", 9)));
ASSERT_EQ(between_code, -3);
}

TEST_F(ChunkedDictionaryTest, TestFindByCodeInvalidNegative) {
Expand Down
26 changes: 24 additions & 2 deletions src/paimon/common/file_index/rangebitmap/dictionary/key_factory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
#include "fmt/format.h"
#include "paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary.h"
#include "paimon/common/file_index/rangebitmap/dictionary/fixed_length_chunk.h"
#include "paimon/common/file_index/rangebitmap/dictionary/variable_length_chunk.h"
#include "paimon/common/file_index/rangebitmap/utils/literal_serialization_utils.h"
#include "paimon/common/utils/field_type_utils.h"
#include "paimon/common/utils/fields_comparator.h"
Expand All @@ -48,6 +49,8 @@ Result<std::shared_ptr<KeyFactory>> KeyFactory::Create(FieldType field_type) {
return std::make_shared<FloatKeyFactory>();
case FieldType::DOUBLE:
return std::make_shared<DoubleKeyFactory>();
case FieldType::STRING:
return std::make_shared<StringKeyFactory>();
default:
return Status::Invalid(fmt::format("Unsupported field type for KeyFactory: {}",
FieldTypeUtils::FieldTypeToString(field_type)));
Expand Down Expand Up @@ -91,12 +94,31 @@ Result<std::unique_ptr<Chunk>> FixedLengthKeyFactory::MmapChunk(
Result<std::unique_ptr<Chunk>> VariableLengthKeyFactory::CreateChunk(
const Literal& key, int32_t code, int32_t keys_length_limit,
const std::shared_ptr<MemoryPool>& pool) {
return Status::NotImplemented("VariableLengthKeyFactory::CreateChunk not implemented");
PAIMON_ASSIGN_OR_RAISE(LiteralSerDeUtils::Serializer serializer,
LiteralSerDeUtils::CreateValueWriter(GetFieldType()));
return std::make_unique<VariableLengthChunk>(key, code, keys_length_limit, shared_from_this(),
serializer, pool);
}
Result<std::unique_ptr<Chunk>> VariableLengthKeyFactory::MmapChunk(
const std::shared_ptr<InputStream>& input_stream, int32_t chunk_offset,
int32_t keys_base_offset, const std::shared_ptr<MemoryPool>& pool) {
return Status::NotImplemented("VariableLengthKeyFactory::MmapChunk not implemented");
PAIMON_RETURN_NOT_OK(input_stream->Seek(chunk_offset, FS_SEEK_SET));
const auto data_in = std::make_shared<DataInputStream>(input_stream);
PAIMON_ASSIGN_OR_RAISE(int8_t version, data_in->ReadValue<int8_t>());
if (version != VariableLengthChunk::kCurrentVersion) {
return Status::Invalid(fmt::format("Unsupported version for KeyFactory: {}", version));
}
PAIMON_ASSIGN_OR_RAISE(LiteralSerDeUtils::Deserializer deserializer,
LiteralSerDeUtils::CreateValueReader(GetFieldType()));
PAIMON_ASSIGN_OR_RAISE(Literal key_literal, deserializer(data_in, pool.get()));
PAIMON_ASSIGN_OR_RAISE(int32_t code, data_in->ReadValue<int32_t>());
PAIMON_ASSIGN_OR_RAISE(int32_t offset, data_in->ReadValue<int32_t>());
PAIMON_ASSIGN_OR_RAISE(int32_t size, data_in->ReadValue<int32_t>());
PAIMON_ASSIGN_OR_RAISE(int32_t offsets_length, data_in->ReadValue<int32_t>());
PAIMON_ASSIGN_OR_RAISE(int32_t keys_length, data_in->ReadValue<int32_t>());
return std::make_unique<VariableLengthChunk>(key_literal, code, offset, size,
shared_from_this(), input_stream, keys_base_offset,
offsets_length, keys_length, pool);
}

/// Java-compatible ordering for floats
Expand Down
Loading
Loading