mirror of
https://github.com/quickwit-oss/tantivy.git
synced 2026-01-10 19:12:54 +00:00
Compare commits
10 Commits
merge-prop
...
columnar-m
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e9af721e60 | ||
|
|
88ed3d8b48 | ||
|
|
aa8408a979 | ||
|
|
4319d8c1bd | ||
|
|
33d18d0424 | ||
|
|
3de018c49f | ||
|
|
96485f21d6 | ||
|
|
1330e6f10d | ||
|
|
5086914304 | ||
|
|
d7a8053cc2 |
6
.github/workflows/coverage.yml
vendored
6
.github/workflows/coverage.yml
vendored
@@ -2,9 +2,9 @@ name: Coverage
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
branches: [ main ]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
branches: [ main ]
|
||||
|
||||
jobs:
|
||||
coverage:
|
||||
@@ -16,7 +16,7 @@ jobs:
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- uses: taiki-e/install-action@cargo-llvm-cov
|
||||
- name: Generate code coverage
|
||||
run: cargo +nightly llvm-cov --all-features --workspace --doctests --lcov --output-path lcov.info
|
||||
run: cargo +nightly llvm-cov --all-features --workspace --lcov --output-path lcov.info
|
||||
- name: Upload coverage to Codecov
|
||||
uses: codecov/codecov-action@v3
|
||||
continue-on-error: true
|
||||
|
||||
2
.gitignore
vendored
2
.gitignore
vendored
@@ -13,5 +13,3 @@ benchmark
|
||||
.idea
|
||||
trace.dat
|
||||
cargo-timing*
|
||||
control
|
||||
variable
|
||||
|
||||
19
Cargo.toml
19
Cargo.toml
@@ -23,7 +23,7 @@ regex = { version = "1.5.5", default-features = false, features = ["std", "unico
|
||||
aho-corasick = "0.7"
|
||||
tantivy-fst = "0.4.0"
|
||||
memmap2 = { version = "0.5.3", optional = true }
|
||||
lz4_flex = { version = "0.10", default-features = false, features = ["checked-decode"], optional = true }
|
||||
lz4_flex = { version = "0.9.2", default-features = false, features = ["checked-decode"], optional = true }
|
||||
brotli = { version = "3.3.4", optional = true }
|
||||
zstd = { version = "0.12", optional = true, default-features = false }
|
||||
snap = { version = "1.0.5", optional = true }
|
||||
@@ -32,7 +32,7 @@ log = "0.4.16"
|
||||
serde = { version = "1.0.136", features = ["derive"] }
|
||||
serde_json = "1.0.79"
|
||||
num_cpus = "1.13.1"
|
||||
fs4 = { version = "0.6.3", optional = true }
|
||||
fs2 = { version = "0.4.3", optional = true }
|
||||
levenshtein_automata = "0.2.1"
|
||||
uuid = { version = "1.0.0", features = ["v4", "serde"] }
|
||||
crossbeam-channel = "0.5.4"
|
||||
@@ -44,23 +44,23 @@ rustc-hash = "1.1.0"
|
||||
thiserror = "1.0.30"
|
||||
htmlescape = "0.3.1"
|
||||
fail = "0.5.0"
|
||||
murmurhash32 = "0.3.0"
|
||||
murmurhash32 = "0.2.0"
|
||||
time = { version = "0.3.10", features = ["serde-well-known"] }
|
||||
smallvec = "1.8.0"
|
||||
rayon = "1.5.2"
|
||||
lru = "0.10.0"
|
||||
lru = "0.9.0"
|
||||
fastdivide = "0.4.0"
|
||||
itertools = "0.10.3"
|
||||
measure_time = "0.8.2"
|
||||
async-trait = "0.1.53"
|
||||
arc-swap = "1.5.0"
|
||||
|
||||
columnar = { version="0.1", path="./columnar", package ="tantivy-columnar" }
|
||||
sstable = { version="0.1", path="./sstable", package ="tantivy-sstable", optional = true }
|
||||
stacker = { version="0.1", path="./stacker", package ="tantivy-stacker" }
|
||||
query-grammar = { version= "0.19.0", path="./query-grammar", package = "tantivy-query-grammar" }
|
||||
tantivy-bitpacker = { version= "0.3", path="./bitpacker" }
|
||||
common = { version= "0.5", path = "./common/", package = "tantivy-common" }
|
||||
tantivy-query-grammar = { version= "0.19.0", path="./query-grammar" }
|
||||
tantivy-bitpacker = { version= "0.3", path="./bitpacker" }
|
||||
columnar = { version= "0.1", path="./columnar", package="tantivy-columnar" }
|
||||
common = { version= "0.5", path = "./common/", package = "tantivy-common" }
|
||||
tokenizer-api = { version="0.1", path="./tokenizer-api", package="tantivy-tokenizer-api" }
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
@@ -77,7 +77,6 @@ test-log = "0.2.10"
|
||||
env_logger = "0.10.0"
|
||||
pprof = { version = "0.11.0", features = ["flamegraph", "criterion"] }
|
||||
futures = "0.3.21"
|
||||
paste = "1.0.11"
|
||||
|
||||
[dev-dependencies.fail]
|
||||
version = "0.5.0"
|
||||
@@ -94,7 +93,7 @@ overflow-checks = true
|
||||
|
||||
[features]
|
||||
default = ["mmap", "stopwords", "lz4-compression"]
|
||||
mmap = ["fs4", "tempfile", "memmap2"]
|
||||
mmap = ["fs2", "tempfile", "memmap2"]
|
||||
stopwords = []
|
||||
|
||||
brotli-compression = ["brotli"]
|
||||
|
||||
23
ci/before_deploy.ps1
Normal file
23
ci/before_deploy.ps1
Normal file
@@ -0,0 +1,23 @@
|
||||
# This script takes care of packaging the build artifacts that will go in the
|
||||
# release zipfile
|
||||
|
||||
$SRC_DIR = $PWD.Path
|
||||
$STAGE = [System.Guid]::NewGuid().ToString()
|
||||
|
||||
Set-Location $ENV:Temp
|
||||
New-Item -Type Directory -Name $STAGE
|
||||
Set-Location $STAGE
|
||||
|
||||
$ZIP = "$SRC_DIR\$($Env:CRATE_NAME)-$($Env:APPVEYOR_REPO_TAG_NAME)-$($Env:TARGET).zip"
|
||||
|
||||
# TODO Update this to package the right artifacts
|
||||
Copy-Item "$SRC_DIR\target\$($Env:TARGET)\release\hello.exe" '.\'
|
||||
|
||||
7z a "$ZIP" *
|
||||
|
||||
Push-AppveyorArtifact "$ZIP"
|
||||
|
||||
Remove-Item *.* -Force
|
||||
Set-Location ..
|
||||
Remove-Item $STAGE
|
||||
Set-Location $SRC_DIR
|
||||
33
ci/before_deploy.sh
Normal file
33
ci/before_deploy.sh
Normal file
@@ -0,0 +1,33 @@
|
||||
# This script takes care of building your crate and packaging it for release
|
||||
|
||||
set -ex
|
||||
|
||||
main() {
|
||||
local src=$(pwd) \
|
||||
stage=
|
||||
|
||||
case $TRAVIS_OS_NAME in
|
||||
linux)
|
||||
stage=$(mktemp -d)
|
||||
;;
|
||||
osx)
|
||||
stage=$(mktemp -d -t tmp)
|
||||
;;
|
||||
esac
|
||||
|
||||
test -f Cargo.lock || cargo generate-lockfile
|
||||
|
||||
# TODO Update this to build the artifacts that matter to you
|
||||
cross rustc --bin hello --target $TARGET --release -- -C lto
|
||||
|
||||
# TODO Update this to package the right artifacts
|
||||
cp target/$TARGET/release/hello $stage/
|
||||
|
||||
cd $stage
|
||||
tar czf $src/$CRATE_NAME-$TRAVIS_TAG-$TARGET.tar.gz *
|
||||
cd $src
|
||||
|
||||
rm -rf $stage
|
||||
}
|
||||
|
||||
main
|
||||
47
ci/install.sh
Normal file
47
ci/install.sh
Normal file
@@ -0,0 +1,47 @@
|
||||
set -ex
|
||||
|
||||
main() {
|
||||
local target=
|
||||
if [ $TRAVIS_OS_NAME = linux ]; then
|
||||
target=x86_64-unknown-linux-musl
|
||||
sort=sort
|
||||
else
|
||||
target=x86_64-apple-darwin
|
||||
sort=gsort # for `sort --sort-version`, from brew's coreutils.
|
||||
fi
|
||||
|
||||
# Builds for iOS are done on OSX, but require the specific target to be
|
||||
# installed.
|
||||
case $TARGET in
|
||||
aarch64-apple-ios)
|
||||
rustup target install aarch64-apple-ios
|
||||
;;
|
||||
armv7-apple-ios)
|
||||
rustup target install armv7-apple-ios
|
||||
;;
|
||||
armv7s-apple-ios)
|
||||
rustup target install armv7s-apple-ios
|
||||
;;
|
||||
i386-apple-ios)
|
||||
rustup target install i386-apple-ios
|
||||
;;
|
||||
x86_64-apple-ios)
|
||||
rustup target install x86_64-apple-ios
|
||||
;;
|
||||
esac
|
||||
|
||||
# This fetches latest stable release
|
||||
local tag=$(git ls-remote --tags --refs --exit-code https://github.com/japaric/cross \
|
||||
| cut -d/ -f3 \
|
||||
| grep -E '^v[0.1.0-9.]+$' \
|
||||
| $sort --version-sort \
|
||||
| tail -n1)
|
||||
curl -LSfs https://japaric.github.io/trust/install.sh | \
|
||||
sh -s -- \
|
||||
--force \
|
||||
--git japaric/cross \
|
||||
--tag $tag \
|
||||
--target $target
|
||||
}
|
||||
|
||||
main
|
||||
30
ci/script.sh
Normal file
30
ci/script.sh
Normal file
@@ -0,0 +1,30 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# This script takes care of testing your crate
|
||||
|
||||
set -ex
|
||||
|
||||
main() {
|
||||
if [ ! -z $CODECOV ]; then
|
||||
echo "Codecov"
|
||||
cargo build --verbose && cargo coverage --verbose --all && bash <(curl -s https://codecov.io/bash) -s target/kcov
|
||||
else
|
||||
echo "Build"
|
||||
cross build --target $TARGET
|
||||
if [ ! -z $DISABLE_TESTS ]; then
|
||||
return
|
||||
fi
|
||||
echo "Test"
|
||||
cross test --target $TARGET --no-default-features --features mmap
|
||||
cross test --target $TARGET --no-default-features --features mmap query-grammar
|
||||
fi
|
||||
for example in $(ls examples/*.rs)
|
||||
do
|
||||
cargo run --example $(basename $example .rs)
|
||||
done
|
||||
}
|
||||
|
||||
# we don't run the "test phase" when doing deploys
|
||||
if [ -z $TRAVIS_TAG ]; then
|
||||
main
|
||||
fi
|
||||
@@ -5,25 +5,24 @@ edition = "2021"
|
||||
license = "MIT"
|
||||
|
||||
[dependencies]
|
||||
itertools = "0.10.5"
|
||||
log = "0.4.17"
|
||||
fnv = "1.0.7"
|
||||
fastdivide = "0.4.0"
|
||||
rand = { version = "0.8.5", optional = true }
|
||||
measure_time = { version = "0.8.2", optional = true }
|
||||
prettytable-rs = { version = "0.10.0", optional = true }
|
||||
|
||||
stacker = { path = "../stacker", package="tantivy-stacker"}
|
||||
serde_json = "1"
|
||||
thiserror = "1"
|
||||
fnv = "1"
|
||||
sstable = { path = "../sstable", package = "tantivy-sstable" }
|
||||
common = { path = "../common", package = "tantivy-common" }
|
||||
itertools = "0.10"
|
||||
log = "0.4"
|
||||
tantivy-bitpacker = { version= "0.3", path = "../bitpacker/" }
|
||||
serde = "1.0.152"
|
||||
prettytable-rs = {version="0.10.0", optional= true}
|
||||
rand = {version="0.8.3", optional= true}
|
||||
fastdivide = "0.4"
|
||||
measure_time = { version="0.8.2", optional=true}
|
||||
|
||||
[dev-dependencies]
|
||||
proptest = "1"
|
||||
more-asserts = "0.3.1"
|
||||
rand = "0.8.5"
|
||||
serde_json = "1"
|
||||
more-asserts = "0.3.0"
|
||||
rand = "0.8.3"
|
||||
|
||||
[features]
|
||||
unstable = []
|
||||
|
||||
6
columnar/Makefile
Normal file
6
columnar/Makefile
Normal file
@@ -0,0 +1,6 @@
|
||||
test:
|
||||
echo "Run test only... No examples."
|
||||
cargo test --tests --lib
|
||||
|
||||
fmt:
|
||||
cargo +nightly fmt --all
|
||||
@@ -28,7 +28,12 @@ fn get_u128_column_random() -> Arc<dyn ColumnValues<u128>> {
|
||||
|
||||
fn get_u128_column_from_data(data: &[u128]) -> Arc<dyn ColumnValues<u128>> {
|
||||
let mut out = vec![];
|
||||
tantivy_columnar::column_values::serialize_column_values_u128(&data, &mut out).unwrap();
|
||||
tantivy_columnar::column_values::serialize_column_values_u128(
|
||||
&(|| data.iter().copied()),
|
||||
data.len() as u32,
|
||||
&mut out,
|
||||
)
|
||||
.unwrap();
|
||||
let out = OwnedBytes::new(out);
|
||||
tantivy_columnar::column_values::open_u128_mapped::<u128>(out).unwrap()
|
||||
}
|
||||
@@ -36,7 +41,7 @@ fn get_u128_column_from_data(data: &[u128]) -> Arc<dyn ColumnValues<u128>> {
|
||||
const FIFTY_PERCENT_RANGE: RangeInclusive<u64> = 1..=50;
|
||||
const SINGLE_ITEM: u64 = 90;
|
||||
const SINGLE_ITEM_RANGE: RangeInclusive<u64> = 90..=90;
|
||||
|
||||
const ONE_PERCENT_ITEM_RANGE: RangeInclusive<u64> = 49..=49;
|
||||
fn get_data_50percent_item() -> Vec<u128> {
|
||||
let mut rng = StdRng::from_seed([1u8; 32]);
|
||||
|
||||
@@ -58,7 +63,7 @@ fn bench_intfastfield_getrange_u128_50percent_hit(b: &mut Bencher) {
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(
|
||||
column.get_docids_for_value_range(
|
||||
*FIFTY_PERCENT_RANGE.start() as u128..=*FIFTY_PERCENT_RANGE.end() as u128,
|
||||
0..data.len() as u32,
|
||||
&mut positions,
|
||||
@@ -74,7 +79,7 @@ fn bench_intfastfield_getrange_u128_single_hit(b: &mut Bencher) {
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(
|
||||
column.get_docids_for_value_range(
|
||||
*SINGLE_ITEM_RANGE.start() as u128..=*SINGLE_ITEM_RANGE.end() as u128,
|
||||
0..data.len() as u32,
|
||||
&mut positions,
|
||||
@@ -90,7 +95,7 @@ fn bench_intfastfield_getrange_u128_hit_all(b: &mut Bencher) {
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(0..=u128::MAX, 0..data.len() as u32, &mut positions);
|
||||
column.get_docids_for_value_range(0..=u128::MAX, 0..data.len() as u32, &mut positions);
|
||||
positions
|
||||
});
|
||||
}
|
||||
|
||||
@@ -5,7 +5,9 @@ use std::ops::RangeInclusive;
|
||||
use std::sync::Arc;
|
||||
|
||||
use rand::prelude::*;
|
||||
use tantivy_columnar::column_values::{serialize_and_load_u64_based_column_values, CodecType};
|
||||
use tantivy_columnar::column_values::{
|
||||
serialize_and_load_u64_based_column_values, CodecType, ALL_U64_CODEC_TYPES,
|
||||
};
|
||||
use tantivy_columnar::*;
|
||||
use test::Bencher;
|
||||
|
||||
@@ -89,7 +91,7 @@ fn bench_intfastfield_getrange_u64_50percent_hit(b: &mut Bencher) {
|
||||
let column: Arc<dyn ColumnValues<u64>> = serialize_and_load(&data, CodecType::Bitpacked);
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(
|
||||
column.get_docids_for_value_range(
|
||||
FIFTY_PERCENT_RANGE,
|
||||
0..data.len() as u32,
|
||||
&mut positions,
|
||||
@@ -106,7 +108,7 @@ fn bench_intfastfield_getrange_u64_1percent_hit(b: &mut Bencher) {
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(
|
||||
column.get_docids_for_value_range(
|
||||
ONE_PERCENT_ITEM_RANGE,
|
||||
0..data.len() as u32,
|
||||
&mut positions,
|
||||
@@ -123,7 +125,7 @@ fn bench_intfastfield_getrange_u64_single_hit(b: &mut Bencher) {
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(SINGLE_ITEM_RANGE, 0..data.len() as u32, &mut positions);
|
||||
column.get_docids_for_value_range(SINGLE_ITEM_RANGE, 0..data.len() as u32, &mut positions);
|
||||
positions
|
||||
});
|
||||
}
|
||||
@@ -136,7 +138,7 @@ fn bench_intfastfield_getrange_u64_hit_all(b: &mut Bencher) {
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(0..=u64::MAX, 0..data.len() as u32, &mut positions);
|
||||
column.get_docids_for_value_range(0..=u64::MAX, 0..data.len() as u32, &mut positions);
|
||||
positions
|
||||
});
|
||||
}
|
||||
@@ -178,7 +180,7 @@ fn bench_intfastfield_scan_all_fflookup(b: &mut Bencher) {
|
||||
b.iter(|| {
|
||||
let mut a = 0u64;
|
||||
for i in 0u32..n as u32 {
|
||||
a += column_ref.get_val(i);
|
||||
a += column.get_val(i);
|
||||
}
|
||||
a
|
||||
});
|
||||
|
||||
@@ -1,17 +0,0 @@
|
||||
[package]
|
||||
name = "tantivy-columnar-cli"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
|
||||
[dependencies]
|
||||
columnar = {path="../", package="tantivy-columnar"}
|
||||
serde_json = "1"
|
||||
serde_json_borrow = {git="https://github.com/PSeitz/serde_json_borrow/"}
|
||||
serde = "1"
|
||||
|
||||
[workspace]
|
||||
members = []
|
||||
|
||||
[profile.release]
|
||||
debug = true
|
||||
@@ -1,134 +0,0 @@
|
||||
use columnar::ColumnarWriter;
|
||||
use columnar::NumericalValue;
|
||||
use serde_json_borrow;
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::io::BufRead;
|
||||
use std::io::BufReader;
|
||||
use std::time::Instant;
|
||||
|
||||
#[derive(Default)]
|
||||
struct JsonStack {
|
||||
path: String,
|
||||
stack: Vec<usize>,
|
||||
}
|
||||
|
||||
impl JsonStack {
|
||||
fn push(&mut self, seg: &str) {
|
||||
let len = self.path.len();
|
||||
self.stack.push(len);
|
||||
self.path.push('.');
|
||||
self.path.push_str(seg);
|
||||
}
|
||||
|
||||
fn pop(&mut self) {
|
||||
if let Some(len) = self.stack.pop() {
|
||||
self.path.truncate(len);
|
||||
}
|
||||
}
|
||||
|
||||
fn path(&self) -> &str {
|
||||
&self.path[1..]
|
||||
}
|
||||
}
|
||||
|
||||
fn append_json_to_columnar(
|
||||
doc: u32,
|
||||
json_value: &serde_json_borrow::Value,
|
||||
columnar: &mut ColumnarWriter,
|
||||
stack: &mut JsonStack,
|
||||
) -> usize {
|
||||
let mut count = 0;
|
||||
match json_value {
|
||||
serde_json_borrow::Value::Null => {}
|
||||
serde_json_borrow::Value::Bool(val) => {
|
||||
columnar.record_numerical(
|
||||
doc,
|
||||
stack.path(),
|
||||
NumericalValue::from(if *val { 1u64 } else { 0u64 }),
|
||||
);
|
||||
count += 1;
|
||||
}
|
||||
serde_json_borrow::Value::Number(num) => {
|
||||
let numerical_value: NumericalValue = if let Some(num_i64) = num.as_i64() {
|
||||
num_i64.into()
|
||||
} else if let Some(num_u64) = num.as_u64() {
|
||||
num_u64.into()
|
||||
} else if let Some(num_f64) = num.as_f64() {
|
||||
num_f64.into()
|
||||
} else {
|
||||
panic!();
|
||||
};
|
||||
count += 1;
|
||||
columnar.record_numerical(
|
||||
doc,
|
||||
stack.path(),
|
||||
numerical_value,
|
||||
);
|
||||
}
|
||||
serde_json_borrow::Value::Str(msg) => {
|
||||
columnar.record_str(
|
||||
doc,
|
||||
stack.path(),
|
||||
msg,
|
||||
);
|
||||
count += 1;
|
||||
},
|
||||
serde_json_borrow::Value::Array(vals) => {
|
||||
for val in vals {
|
||||
count += append_json_to_columnar(doc, val, columnar, stack);
|
||||
}
|
||||
},
|
||||
serde_json_borrow::Value::Object(json_map) => {
|
||||
for (child_key, child_val) in json_map {
|
||||
stack.push(child_key);
|
||||
count += append_json_to_columnar(doc, child_val, columnar, stack);
|
||||
stack.pop();
|
||||
}
|
||||
},
|
||||
}
|
||||
count
|
||||
}
|
||||
|
||||
fn main() -> io::Result<()> {
|
||||
let file = File::open("gh_small.json")?;
|
||||
let mut reader = BufReader::new(file);
|
||||
let mut line = String::with_capacity(100);
|
||||
let mut columnar = columnar::ColumnarWriter::default();
|
||||
let mut doc = 0;
|
||||
let start = Instant::now();
|
||||
let mut stack = JsonStack::default();
|
||||
let mut total_count = 0;
|
||||
|
||||
let start_build = Instant::now();
|
||||
loop {
|
||||
line.clear();
|
||||
let len = reader.read_line(&mut line)?;
|
||||
if len == 0 {
|
||||
break;
|
||||
}
|
||||
let Ok(json_value) = serde_json::from_str::<serde_json_borrow::Value>(&line) else { continue; };
|
||||
total_count += append_json_to_columnar(doc, &json_value, &mut columnar, &mut stack);
|
||||
doc += 1;
|
||||
}
|
||||
println!("Build in {:?}", start_build.elapsed());
|
||||
|
||||
println!("value count {total_count}");
|
||||
|
||||
let mut buffer = Vec::new();
|
||||
let start_serialize = Instant::now();
|
||||
columnar.serialize(doc, None, &mut buffer)?;
|
||||
println!("Serialized in {:?}", start_serialize.elapsed());
|
||||
println!("num docs: {doc}, {:?}", start.elapsed());
|
||||
println!("buffer len {} MB", buffer.len() / 1_000_000);
|
||||
let columnar = columnar::ColumnarReader::open(buffer)?;
|
||||
for (column_name, dynamic_column) in columnar.list_columns()? {
|
||||
let num_bytes = dynamic_column.num_bytes();
|
||||
let typ = dynamic_column.column_type();
|
||||
if num_bytes > 1_000_000 {
|
||||
println!("{column_name} {typ:?} {} KB", num_bytes / 1_000);
|
||||
}
|
||||
}
|
||||
println!("{} columns", columnar.num_columns());
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,16 +1,14 @@
|
||||
# zero to one
|
||||
|
||||
* merges with non trivial mapping (deletes / sort)
|
||||
* emission of the sort mapping.
|
||||
+ muttivaued range queries restrat frm the beginning all of the time.
|
||||
* revisit line codec
|
||||
* add columns from schema on merge
|
||||
* removal of all rows of a column in the schema due to deletes
|
||||
* Plugging JSON
|
||||
* replug examples
|
||||
* move datetime to quickwit common
|
||||
* switch to nanos
|
||||
* reintroduce the gcd map.
|
||||
replug examples
|
||||
|
||||
|
||||
# Perf and Size
|
||||
* remove alloc in `ord_to_term`
|
||||
+ multivaued range queries restrat frm the beginning all of the time.
|
||||
* re-add ZSTD compression for dictionaries
|
||||
no systematic monotonic mapping
|
||||
consider removing multilinear
|
||||
|
||||
@@ -1,36 +0,0 @@
|
||||
use crate::{Column, DocId, RowId};
|
||||
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct ColumnBlockAccessor<T> {
|
||||
val_cache: Vec<T>,
|
||||
docid_cache: Vec<DocId>,
|
||||
row_id_cache: Vec<RowId>,
|
||||
}
|
||||
|
||||
impl<T: PartialOrd + Copy + std::fmt::Debug + Send + Sync + 'static + Default>
|
||||
ColumnBlockAccessor<T>
|
||||
{
|
||||
#[inline]
|
||||
pub fn fetch_block(&mut self, docs: &[u32], accessor: &Column<T>) {
|
||||
self.docid_cache.clear();
|
||||
self.row_id_cache.clear();
|
||||
accessor.row_ids_for_docs(docs, &mut self.docid_cache, &mut self.row_id_cache);
|
||||
self.val_cache.resize(self.row_id_cache.len(), T::default());
|
||||
accessor
|
||||
.values
|
||||
.get_vals(&self.row_id_cache, &mut self.val_cache);
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn iter_vals(&self) -> impl Iterator<Item = T> + '_ {
|
||||
self.val_cache.iter().cloned()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn iter_docid_vals(&self) -> impl Iterator<Item = (DocId, T)> + '_ {
|
||||
self.docid_cache
|
||||
.iter()
|
||||
.cloned()
|
||||
.zip(self.val_cache.iter().cloned())
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
use std::io;
|
||||
use std::ops::Deref;
|
||||
use std::sync::Arc;
|
||||
use std::{fmt, io};
|
||||
|
||||
use sstable::{Dictionary, VoidSSTable};
|
||||
|
||||
@@ -21,14 +21,6 @@ pub struct BytesColumn {
|
||||
pub(crate) term_ord_column: Column<u64>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for BytesColumn {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("BytesColumn")
|
||||
.field("term_ord_column", &self.term_ord_column)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl BytesColumn {
|
||||
/// Fills the given `output` buffer with the term associated to the ordinal `ord`.
|
||||
///
|
||||
@@ -40,11 +32,11 @@ impl BytesColumn {
|
||||
|
||||
/// Returns the number of rows in the column.
|
||||
pub fn num_rows(&self) -> RowId {
|
||||
self.term_ord_column.num_docs()
|
||||
self.term_ord_column.num_rows()
|
||||
}
|
||||
|
||||
pub fn term_ords(&self, row_id: RowId) -> impl Iterator<Item = u64> + '_ {
|
||||
self.term_ord_column.values_for_doc(row_id)
|
||||
self.term_ord_column.values(row_id)
|
||||
}
|
||||
|
||||
/// Returns the column of ordinals
|
||||
@@ -64,9 +56,9 @@ impl BytesColumn {
|
||||
#[derive(Clone)]
|
||||
pub struct StrColumn(BytesColumn);
|
||||
|
||||
impl fmt::Debug for StrColumn {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
write!(f, "{:?}", self.term_ord_column)
|
||||
impl From<BytesColumn> for StrColumn {
|
||||
fn from(bytes_col: BytesColumn) -> Self {
|
||||
StrColumn(bytes_col)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -77,10 +69,6 @@ impl From<StrColumn> for BytesColumn {
|
||||
}
|
||||
|
||||
impl StrColumn {
|
||||
pub(crate) fn wrap(bytes_column: BytesColumn) -> StrColumn {
|
||||
StrColumn(bytes_column)
|
||||
}
|
||||
|
||||
pub fn dictionary(&self) -> &Dictionary<VoidSSTable> {
|
||||
self.0.dictionary.as_ref()
|
||||
}
|
||||
|
||||
@@ -1,49 +1,29 @@
|
||||
mod dictionary_encoded;
|
||||
mod serialize;
|
||||
|
||||
use core::fmt;
|
||||
use std::fmt::Debug;
|
||||
use std::io::Write;
|
||||
use std::ops::{Deref, Range, RangeInclusive};
|
||||
use std::ops::Deref;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::BinarySerializable;
|
||||
pub use dictionary_encoded::{BytesColumn, StrColumn};
|
||||
pub use serialize::{
|
||||
open_column_bytes, open_column_str, open_column_u128, open_column_u64,
|
||||
serialize_column_mappable_to_u128, serialize_column_mappable_to_u64,
|
||||
open_column_bytes, open_column_u128, open_column_u64, serialize_column_mappable_to_u128,
|
||||
serialize_column_mappable_to_u64,
|
||||
};
|
||||
|
||||
use crate::column_index::ColumnIndex;
|
||||
use crate::column_values::monotonic_mapping::StrictlyMonotonicMappingToInternal;
|
||||
use crate::column_values::{monotonic_map_column, ColumnValues};
|
||||
use crate::{Cardinality, DocId, EmptyColumnValues, MonotonicallyMappableToU64, RowId};
|
||||
use crate::{Cardinality, MonotonicallyMappableToU64, RowId};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Column<T = u64> {
|
||||
pub index: ColumnIndex,
|
||||
pub idx: ColumnIndex,
|
||||
pub values: Arc<dyn ColumnValues<T>>,
|
||||
}
|
||||
|
||||
impl<T: fmt::Debug + PartialOrd + Send + Sync + Copy + 'static> fmt::Debug for Column<T> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
let num_docs = self.num_docs();
|
||||
let entries = (0..num_docs)
|
||||
.map(|i| (i, self.values_for_doc(i).collect::<Vec<_>>()))
|
||||
.filter(|(_, vals)| !vals.is_empty());
|
||||
f.debug_map().entries(entries).finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: PartialOrd + Default> Column<T> {
|
||||
pub fn build_empty_column(num_docs: u32) -> Column<T> {
|
||||
Column {
|
||||
index: ColumnIndex::Empty { num_docs },
|
||||
values: Arc::new(EmptyColumnValues),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: MonotonicallyMappableToU64> Column<T> {
|
||||
pub fn to_u64_monotonic(self) -> Column<u64> {
|
||||
let values = Arc::new(monotonic_map_column(
|
||||
@@ -51,27 +31,24 @@ impl<T: MonotonicallyMappableToU64> Column<T> {
|
||||
StrictlyMonotonicMappingToInternal::<T>::new(),
|
||||
));
|
||||
Column {
|
||||
index: self.index,
|
||||
idx: self.idx,
|
||||
values,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: PartialOrd + Copy + Debug + Send + Sync + 'static> Column<T> {
|
||||
#[inline]
|
||||
pub fn get_cardinality(&self) -> Cardinality {
|
||||
self.index.get_cardinality()
|
||||
self.idx.get_cardinality()
|
||||
}
|
||||
|
||||
pub fn num_docs(&self) -> RowId {
|
||||
match &self.index {
|
||||
ColumnIndex::Empty { num_docs } => *num_docs,
|
||||
ColumnIndex::Full => self.values.num_vals(),
|
||||
ColumnIndex::Optional(optional_index) => optional_index.num_docs(),
|
||||
pub fn num_rows(&self) -> RowId {
|
||||
match &self.idx {
|
||||
ColumnIndex::Full => self.values.num_vals() as u32,
|
||||
ColumnIndex::Optional(optional_index) => optional_index.num_rows(),
|
||||
ColumnIndex::Multivalued(col_index) => {
|
||||
// The multivalued index contains all value start row_id,
|
||||
// and one extra value at the end with the overall number of rows.
|
||||
col_index.num_docs()
|
||||
col_index.num_rows()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -85,61 +62,14 @@ impl<T: PartialOrd + Copy + Debug + Send + Sync + 'static> Column<T> {
|
||||
}
|
||||
|
||||
pub fn first(&self, row_id: RowId) -> Option<T> {
|
||||
self.values_for_doc(row_id).next()
|
||||
self.values(row_id).next()
|
||||
}
|
||||
|
||||
/// Translates a block of docis to row_ids.
|
||||
///
|
||||
/// returns the row_ids and the matching docids on the same index
|
||||
/// e.g.
|
||||
/// DocId In: [0, 5, 6]
|
||||
/// DocId Out: [0, 0, 6, 6]
|
||||
/// RowId Out: [0, 1, 2, 3]
|
||||
#[inline]
|
||||
pub fn row_ids_for_docs(
|
||||
&self,
|
||||
doc_ids: &[DocId],
|
||||
doc_ids_out: &mut Vec<DocId>,
|
||||
row_ids: &mut Vec<RowId>,
|
||||
) {
|
||||
self.index.docids_to_rowids(doc_ids, doc_ids_out, row_ids)
|
||||
}
|
||||
|
||||
pub fn values_for_doc(&self, doc_id: DocId) -> impl Iterator<Item = T> + '_ {
|
||||
self.value_row_ids(doc_id)
|
||||
pub fn values(&self, row_id: RowId) -> impl Iterator<Item = T> + '_ {
|
||||
self.value_row_ids(row_id)
|
||||
.map(|value_row_id: RowId| self.values.get_val(value_row_id))
|
||||
}
|
||||
|
||||
/// Get the docids of values which are in the provided value range.
|
||||
#[inline]
|
||||
pub fn get_docids_for_value_range(
|
||||
&self,
|
||||
value_range: RangeInclusive<T>,
|
||||
selected_docid_range: Range<u32>,
|
||||
doc_ids: &mut Vec<u32>,
|
||||
) {
|
||||
// convert passed docid range to row id range
|
||||
let rowid_range = self
|
||||
.index
|
||||
.docid_range_to_rowids(selected_docid_range.clone());
|
||||
|
||||
// Load rows
|
||||
self.values
|
||||
.get_row_ids_for_value_range(value_range, rowid_range, doc_ids);
|
||||
// Convert rows to docids
|
||||
self.index
|
||||
.select_batch_in_place(selected_docid_range.start, doc_ids);
|
||||
}
|
||||
|
||||
/// Fils the output vector with the (possibly multiple values that are associated_with
|
||||
/// `row_id`.
|
||||
///
|
||||
/// This method clears the `output` vector.
|
||||
pub fn fill_vals(&self, row_id: RowId, output: &mut Vec<T>) {
|
||||
output.clear();
|
||||
output.extend(self.values_for_doc(row_id));
|
||||
}
|
||||
|
||||
pub fn first_or_default_col(self, default_value: T) -> Arc<dyn ColumnValues<T>> {
|
||||
Arc::new(FirstValueWithDefault {
|
||||
column: self,
|
||||
@@ -152,7 +82,7 @@ impl<T> Deref for Column<T> {
|
||||
type Target = ColumnIndex;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.index
|
||||
&self.idx
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,11 +120,10 @@ impl<T: PartialOrd + Debug + Send + Sync + Copy + 'static> ColumnValues<T>
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
match &self.column.index {
|
||||
ColumnIndex::Empty { .. } => 0u32,
|
||||
match &self.column.idx {
|
||||
ColumnIndex::Full => self.column.values.num_vals(),
|
||||
ColumnIndex::Optional(optional_idx) => optional_idx.num_docs(),
|
||||
ColumnIndex::Multivalued(multivalue_idx) => multivalue_idx.num_docs(),
|
||||
ColumnIndex::Optional(optional_idx) => optional_idx.num_rows(),
|
||||
ColumnIndex::Multivalued(_) => todo!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use std::fmt::Debug;
|
||||
use std::io;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
@@ -7,25 +8,24 @@ use sstable::Dictionary;
|
||||
|
||||
use crate::column::{BytesColumn, Column};
|
||||
use crate::column_index::{serialize_column_index, SerializableColumnIndex};
|
||||
use crate::column_values::{
|
||||
load_u64_based_column_values, serialize_column_values_u128, serialize_u64_based_column_values,
|
||||
CodecType, MonotonicallyMappableToU128, MonotonicallyMappableToU64,
|
||||
};
|
||||
use crate::column_values::serialize::serialize_column_values_u128;
|
||||
use crate::column_values::u64_based::{serialize_u64_based_column_values, CodecType};
|
||||
use crate::column_values::{MonotonicallyMappableToU128, MonotonicallyMappableToU64};
|
||||
use crate::iterable::Iterable;
|
||||
use crate::StrColumn;
|
||||
|
||||
pub fn serialize_column_mappable_to_u128<T: MonotonicallyMappableToU128>(
|
||||
column_index: SerializableColumnIndex<'_>,
|
||||
iterable: &dyn Iterable<T>,
|
||||
num_vals: u32,
|
||||
output: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
let column_index_num_bytes = serialize_column_index(column_index, output)?;
|
||||
serialize_column_values_u128(iterable, output)?;
|
||||
serialize_column_values_u128(iterable, num_vals, output)?;
|
||||
output.write_all(&column_index_num_bytes.to_le_bytes())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn serialize_column_mappable_to_u64<T: MonotonicallyMappableToU64>(
|
||||
pub fn serialize_column_mappable_to_u64<T: MonotonicallyMappableToU64 + Debug>(
|
||||
column_index: SerializableColumnIndex<'_>,
|
||||
column_values: &impl Iterable<T>,
|
||||
output: &mut impl Write,
|
||||
@@ -50,9 +50,10 @@ pub fn open_column_u64<T: MonotonicallyMappableToU64>(bytes: OwnedBytes) -> io::
|
||||
);
|
||||
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
|
||||
let column_index = crate::column_index::open_column_index(column_index_data)?;
|
||||
let column_values = load_u64_based_column_values(column_values_data)?;
|
||||
let column_values =
|
||||
crate::column_values::u64_based::load_u64_based_column_values(column_values_data)?;
|
||||
Ok(Column {
|
||||
index: column_index,
|
||||
idx: column_index,
|
||||
values: column_values,
|
||||
})
|
||||
}
|
||||
@@ -71,24 +72,20 @@ pub fn open_column_u128<T: MonotonicallyMappableToU128>(
|
||||
let column_index = crate::column_index::open_column_index(column_index_data)?;
|
||||
let column_values = crate::column_values::open_u128_mapped(column_values_data)?;
|
||||
Ok(Column {
|
||||
index: column_index,
|
||||
idx: column_index,
|
||||
values: column_values,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn open_column_bytes(data: OwnedBytes) -> io::Result<BytesColumn> {
|
||||
pub fn open_column_bytes<T: From<BytesColumn>>(data: OwnedBytes) -> io::Result<T> {
|
||||
let (body, dictionary_len_bytes) = data.rsplit(4);
|
||||
let dictionary_len = u32::from_le_bytes(dictionary_len_bytes.as_slice().try_into().unwrap());
|
||||
let (dictionary_bytes, column_bytes) = body.split(dictionary_len as usize);
|
||||
let dictionary = Arc::new(Dictionary::from_bytes(dictionary_bytes)?);
|
||||
let term_ord_column = crate::column::open_column_u64::<u64>(column_bytes)?;
|
||||
Ok(BytesColumn {
|
||||
let bytes_column = BytesColumn {
|
||||
dictionary,
|
||||
term_ord_column,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn open_column_str(data: OwnedBytes) -> io::Result<StrColumn> {
|
||||
let bytes_column = open_column_bytes(data)?;
|
||||
Ok(StrColumn::wrap(bytes_column))
|
||||
};
|
||||
Ok(bytes_column.into())
|
||||
}
|
||||
|
||||
@@ -1,19 +1,29 @@
|
||||
use std::iter;
|
||||
|
||||
use crate::column_index::{SerializableColumnIndex, Set};
|
||||
use crate::column_index::{
|
||||
multivalued_index, serialize_column_index, SerializableColumnIndex, Set,
|
||||
};
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{Cardinality, ColumnIndex, RowId, StackMergeOrder};
|
||||
use crate::{Cardinality, ColumnIndex, MergeRowOrder, RowId, StackMergeOrder};
|
||||
|
||||
/// Simple case:
|
||||
/// The new mapping just consists in stacking the different column indexes.
|
||||
///
|
||||
/// There are no sort nor deletes involved.
|
||||
pub fn merge_column_index_stacked<'a>(
|
||||
columns: &'a [ColumnIndex],
|
||||
cardinality_after_merge: Cardinality,
|
||||
stack_merge_order: &'a StackMergeOrder,
|
||||
fn detect_cardinality(columns: &[Option<ColumnIndex>]) -> Cardinality {
|
||||
columns
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(ColumnIndex::get_cardinality)
|
||||
.max()
|
||||
.unwrap_or(Cardinality::Full)
|
||||
}
|
||||
|
||||
pub fn stack_column_index<'a>(
|
||||
columns: &'a [Option<ColumnIndex>],
|
||||
merge_row_order: &'a MergeRowOrder,
|
||||
) -> SerializableColumnIndex<'a> {
|
||||
match cardinality_after_merge {
|
||||
let MergeRowOrder::Stack(stack_merge_order) = merge_row_order else {
|
||||
panic!("only supporting stacking at the moment.");
|
||||
};
|
||||
let cardinality = detect_cardinality(columns);
|
||||
match cardinality {
|
||||
Cardinality::Full => SerializableColumnIndex::Full,
|
||||
Cardinality::Optional => SerializableColumnIndex::Optional {
|
||||
non_null_row_ids: Box::new(StackedOptionalIndex {
|
||||
@@ -33,7 +43,7 @@ pub fn merge_column_index_stacked<'a>(
|
||||
}
|
||||
|
||||
struct StackedOptionalIndex<'a> {
|
||||
columns: &'a [ColumnIndex],
|
||||
columns: &'a [Option<ColumnIndex>],
|
||||
stack_merge_order: &'a StackMergeOrder,
|
||||
}
|
||||
|
||||
@@ -46,16 +56,16 @@ impl<'a> Iterable<RowId> for StackedOptionalIndex<'a> {
|
||||
.flat_map(|(columnar_id, column_index_opt)| {
|
||||
let columnar_row_range = self.stack_merge_order.columnar_range(columnar_id);
|
||||
let rows_it: Box<dyn Iterator<Item = RowId>> = match column_index_opt {
|
||||
ColumnIndex::Full => Box::new(columnar_row_range),
|
||||
ColumnIndex::Optional(optional_index) => Box::new(
|
||||
Some(ColumnIndex::Full) => Box::new(columnar_row_range),
|
||||
Some(ColumnIndex::Optional(optional_index)) => Box::new(
|
||||
optional_index
|
||||
.iter_rows()
|
||||
.map(move |row_id: RowId| columnar_row_range.start + row_id),
|
||||
.map(move |row_id: RowId| row_id + columnar_row_range.start),
|
||||
),
|
||||
ColumnIndex::Multivalued(_) => {
|
||||
Some(ColumnIndex::Multivalued(_)) => {
|
||||
panic!("No multivalued index is allowed when stacking column index");
|
||||
}
|
||||
ColumnIndex::Empty { .. } => Box::new(std::iter::empty()),
|
||||
None => Box::new(std::iter::empty()),
|
||||
};
|
||||
rows_it
|
||||
}),
|
||||
@@ -65,18 +75,18 @@ impl<'a> Iterable<RowId> for StackedOptionalIndex<'a> {
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct StackedMultivaluedIndex<'a> {
|
||||
columns: &'a [ColumnIndex],
|
||||
columns: &'a [Option<ColumnIndex>],
|
||||
stack_merge_order: &'a StackMergeOrder,
|
||||
}
|
||||
|
||||
fn convert_column_opt_to_multivalued_index<'a>(
|
||||
column_index_opt: &'a ColumnIndex,
|
||||
column_index_opt: Option<&'a ColumnIndex>,
|
||||
num_rows: RowId,
|
||||
) -> Box<dyn Iterator<Item = RowId> + 'a> {
|
||||
match column_index_opt {
|
||||
ColumnIndex::Empty { .. } => Box::new(iter::repeat(0u32).take(num_rows as usize + 1)),
|
||||
ColumnIndex::Full => Box::new(0..num_rows + 1),
|
||||
ColumnIndex::Optional(optional_index) => {
|
||||
None => Box::new(iter::repeat(0u32).take(num_rows as usize + 1)),
|
||||
Some(ColumnIndex::Full) => Box::new(0..num_rows + 1),
|
||||
Some(ColumnIndex::Optional(optional_index)) => {
|
||||
Box::new(
|
||||
(0..num_rows)
|
||||
// TODO optimize
|
||||
@@ -84,7 +94,9 @@ fn convert_column_opt_to_multivalued_index<'a>(
|
||||
.chain(std::iter::once(optional_index.num_non_nulls())),
|
||||
)
|
||||
}
|
||||
ColumnIndex::Multivalued(multivalued_index) => multivalued_index.start_index_column.iter(),
|
||||
Some(ColumnIndex::Multivalued(multivalued_index)) => {
|
||||
multivalued_index.start_index_column.iter()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -93,6 +105,7 @@ impl<'a> Iterable<RowId> for StackedMultivaluedIndex<'a> {
|
||||
let multivalued_indexes =
|
||||
self.columns
|
||||
.iter()
|
||||
.map(Option::as_ref)
|
||||
.enumerate()
|
||||
.map(|(columnar_id, column_opt)| {
|
||||
let num_rows =
|
||||
@@ -128,6 +141,13 @@ fn stack_multivalued_indexes<'a>(
|
||||
}))
|
||||
}
|
||||
|
||||
fn stack_multivalued_index<'a>(
|
||||
columns: &'a [Option<ColumnIndex>],
|
||||
stack_merge_order: &StackMergeOrder,
|
||||
) -> Box<dyn Iterable<RowId> + 'a> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::RowId;
|
||||
@@ -1,134 +0,0 @@
|
||||
mod shuffled;
|
||||
mod stacked;
|
||||
|
||||
use shuffled::merge_column_index_shuffled;
|
||||
use stacked::merge_column_index_stacked;
|
||||
|
||||
use crate::column_index::SerializableColumnIndex;
|
||||
use crate::{Cardinality, ColumnIndex, MergeRowOrder};
|
||||
|
||||
// For simplification, we never have cardinality go down due to deletes.
|
||||
fn detect_cardinality(columns: &[ColumnIndex]) -> Cardinality {
|
||||
columns
|
||||
.iter()
|
||||
.map(ColumnIndex::get_cardinality)
|
||||
.max()
|
||||
.unwrap_or(Cardinality::Full)
|
||||
}
|
||||
|
||||
pub fn merge_column_index<'a>(
|
||||
columns: &'a [ColumnIndex],
|
||||
merge_row_order: &'a MergeRowOrder,
|
||||
) -> SerializableColumnIndex<'a> {
|
||||
// For simplification, we do not try to detect whether the cardinality could be
|
||||
// downgraded thanks to deletes.
|
||||
let cardinality_after_merge = detect_cardinality(columns);
|
||||
match merge_row_order {
|
||||
MergeRowOrder::Stack(stack_merge_order) => {
|
||||
merge_column_index_stacked(columns, cardinality_after_merge, stack_merge_order)
|
||||
}
|
||||
MergeRowOrder::Shuffled(complex_merge_order) => {
|
||||
merge_column_index_shuffled(columns, cardinality_after_merge, complex_merge_order)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TODO actually, the shuffled code path is a bit too general.
|
||||
// In practise, we do not really shuffle everything.
|
||||
// The merge order restricted to a specific column keeps the original row order.
|
||||
//
|
||||
// This may offer some optimization that we have not explored yet.
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::column_index::merge::detect_cardinality;
|
||||
use crate::column_index::multivalued_index::MultiValueIndex;
|
||||
use crate::column_index::{merge_column_index, OptionalIndex, SerializableColumnIndex};
|
||||
use crate::{Cardinality, ColumnIndex, MergeRowOrder, RowAddr, RowId, ShuffleMergeOrder};
|
||||
|
||||
#[test]
|
||||
fn test_detect_cardinality() {
|
||||
assert_eq!(detect_cardinality(&[]), Cardinality::Full);
|
||||
let optional_index: ColumnIndex = OptionalIndex::for_test(1, &[]).into();
|
||||
let multivalued_index: ColumnIndex = MultiValueIndex::for_test(&[0, 1]).into();
|
||||
assert_eq!(
|
||||
detect_cardinality(&[optional_index.clone(), ColumnIndex::Empty { num_docs: 0 }]),
|
||||
Cardinality::Optional
|
||||
);
|
||||
assert_eq!(
|
||||
detect_cardinality(&[optional_index.clone(), ColumnIndex::Full]),
|
||||
Cardinality::Optional
|
||||
);
|
||||
assert_eq!(
|
||||
detect_cardinality(&[
|
||||
multivalued_index.clone(),
|
||||
ColumnIndex::Empty { num_docs: 0 }
|
||||
]),
|
||||
Cardinality::Multivalued
|
||||
);
|
||||
assert_eq!(
|
||||
detect_cardinality(&[multivalued_index.clone(), optional_index.clone()]),
|
||||
Cardinality::Multivalued
|
||||
);
|
||||
assert_eq!(
|
||||
detect_cardinality(&[optional_index, multivalued_index]),
|
||||
Cardinality::Multivalued
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_index_multivalued_sorted() {
|
||||
let column_indexes: Vec<ColumnIndex> = vec![MultiValueIndex::for_test(&[0, 2, 5]).into()];
|
||||
let merge_row_order: MergeRowOrder = ShuffleMergeOrder::for_test(
|
||||
&[2],
|
||||
vec![
|
||||
RowAddr {
|
||||
segment_ord: 0u32,
|
||||
row_id: 1u32,
|
||||
},
|
||||
RowAddr {
|
||||
segment_ord: 0u32,
|
||||
row_id: 0u32,
|
||||
},
|
||||
],
|
||||
)
|
||||
.into();
|
||||
let merged_column_index = merge_column_index(&column_indexes[..], &merge_row_order);
|
||||
let SerializableColumnIndex::Multivalued(start_index_iterable) = merged_column_index
|
||||
else { panic!("Excpected a multivalued index") };
|
||||
let start_indexes: Vec<RowId> = start_index_iterable.boxed_iter().collect();
|
||||
assert_eq!(&start_indexes, &[0, 3, 5]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_index_multivalued_sorted_several_segment() {
|
||||
let column_indexes: Vec<ColumnIndex> = vec![
|
||||
MultiValueIndex::for_test(&[0, 2, 5]).into(),
|
||||
ColumnIndex::Empty { num_docs: 0 },
|
||||
MultiValueIndex::for_test(&[0, 1, 4]).into(),
|
||||
];
|
||||
let merge_row_order: MergeRowOrder = ShuffleMergeOrder::for_test(
|
||||
&[2, 0, 2],
|
||||
vec![
|
||||
RowAddr {
|
||||
segment_ord: 2u32,
|
||||
row_id: 1u32,
|
||||
},
|
||||
RowAddr {
|
||||
segment_ord: 0u32,
|
||||
row_id: 0u32,
|
||||
},
|
||||
RowAddr {
|
||||
segment_ord: 2u32,
|
||||
row_id: 0u32,
|
||||
},
|
||||
],
|
||||
)
|
||||
.into();
|
||||
let merged_column_index = merge_column_index(&column_indexes[..], &merge_row_order);
|
||||
let SerializableColumnIndex::Multivalued(start_index_iterable) = merged_column_index
|
||||
else { panic!("Excpected a multivalued index") };
|
||||
let start_indexes: Vec<RowId> = start_index_iterable.boxed_iter().collect();
|
||||
assert_eq!(&start_indexes, &[0, 3, 5, 6]);
|
||||
}
|
||||
}
|
||||
@@ -1,165 +0,0 @@
|
||||
use std::iter;
|
||||
|
||||
use crate::column_index::{SerializableColumnIndex, Set};
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{Cardinality, ColumnIndex, RowId, ShuffleMergeOrder};
|
||||
|
||||
pub fn merge_column_index_shuffled<'a>(
|
||||
column_indexes: &'a [ColumnIndex],
|
||||
cardinality_after_merge: Cardinality,
|
||||
shuffle_merge_order: &'a ShuffleMergeOrder,
|
||||
) -> SerializableColumnIndex<'a> {
|
||||
match cardinality_after_merge {
|
||||
Cardinality::Full => SerializableColumnIndex::Full,
|
||||
Cardinality::Optional => {
|
||||
let non_null_row_ids =
|
||||
merge_column_index_shuffled_optional(column_indexes, shuffle_merge_order);
|
||||
SerializableColumnIndex::Optional {
|
||||
non_null_row_ids,
|
||||
num_rows: shuffle_merge_order.num_rows(),
|
||||
}
|
||||
}
|
||||
Cardinality::Multivalued => {
|
||||
let multivalue_start_index =
|
||||
merge_column_index_shuffled_multivalued(column_indexes, shuffle_merge_order);
|
||||
SerializableColumnIndex::Multivalued(multivalue_start_index)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Merge several column indexes into one, ordering rows according to the merge_order passed as
|
||||
/// argument. While it is true that the `merge_order` may imply deletes and hence could in theory a
|
||||
/// multivalued index into an optional one, this is not supported today for simplification.
|
||||
///
|
||||
/// In other words the column_indexes passed as argument may NOT be multivalued.
|
||||
fn merge_column_index_shuffled_optional<'a>(
|
||||
column_indexes: &'a [ColumnIndex],
|
||||
merge_order: &'a ShuffleMergeOrder,
|
||||
) -> Box<dyn Iterable<RowId> + 'a> {
|
||||
Box::new(ShuffledIndex {
|
||||
column_indexes,
|
||||
merge_order,
|
||||
})
|
||||
}
|
||||
|
||||
struct ShuffledIndex<'a> {
|
||||
column_indexes: &'a [ColumnIndex],
|
||||
merge_order: &'a ShuffleMergeOrder,
|
||||
}
|
||||
|
||||
impl<'a> Iterable<u32> for ShuffledIndex<'a> {
|
||||
fn boxed_iter(&self) -> Box<dyn Iterator<Item = u32> + '_> {
|
||||
Box::new(
|
||||
self.merge_order
|
||||
.iter_new_to_old_row_addrs()
|
||||
.enumerate()
|
||||
.filter_map(|(new_row_id, old_row_addr)| {
|
||||
let column_index = &self.column_indexes[old_row_addr.segment_ord as usize];
|
||||
let row_id = new_row_id as u32;
|
||||
if column_index.has_value(old_row_addr.row_id) {
|
||||
Some(row_id)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_column_index_shuffled_multivalued<'a>(
|
||||
column_indexes: &'a [ColumnIndex],
|
||||
merge_order: &'a ShuffleMergeOrder,
|
||||
) -> Box<dyn Iterable<RowId> + 'a> {
|
||||
Box::new(ShuffledMultivaluedIndex {
|
||||
column_indexes,
|
||||
merge_order,
|
||||
})
|
||||
}
|
||||
|
||||
struct ShuffledMultivaluedIndex<'a> {
|
||||
column_indexes: &'a [ColumnIndex],
|
||||
merge_order: &'a ShuffleMergeOrder,
|
||||
}
|
||||
|
||||
fn iter_num_values<'a>(
|
||||
column_indexes: &'a [ColumnIndex],
|
||||
merge_order: &'a ShuffleMergeOrder,
|
||||
) -> impl Iterator<Item = u32> + 'a {
|
||||
merge_order.iter_new_to_old_row_addrs().map(|row_addr| {
|
||||
let column_index = &column_indexes[row_addr.segment_ord as usize];
|
||||
match column_index {
|
||||
ColumnIndex::Empty { .. } => 0u32,
|
||||
ColumnIndex::Full => 1,
|
||||
ColumnIndex::Optional(optional_index) => {
|
||||
u32::from(optional_index.contains(row_addr.row_id))
|
||||
}
|
||||
ColumnIndex::Multivalued(multivalued_index) => {
|
||||
multivalued_index.range(row_addr.row_id).len() as u32
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Transforms an iterator containing the number of vals per row (with `num_rows` elements)
|
||||
/// into a `start_offset` iterator starting at 0 and (with `num_rows + 1` element)
|
||||
fn integrate_num_vals(num_vals: impl Iterator<Item = u32>) -> impl Iterator<Item = RowId> {
|
||||
iter::once(0u32).chain(num_vals.scan(0, |state, num_vals| {
|
||||
*state += num_vals;
|
||||
Some(*state)
|
||||
}))
|
||||
}
|
||||
|
||||
impl<'a> Iterable<u32> for ShuffledMultivaluedIndex<'a> {
|
||||
fn boxed_iter(&self) -> Box<dyn Iterator<Item = u32> + '_> {
|
||||
let num_vals_per_row = iter_num_values(self.column_indexes, self.merge_order);
|
||||
Box::new(integrate_num_vals(num_vals_per_row))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::column_index::OptionalIndex;
|
||||
use crate::RowAddr;
|
||||
|
||||
#[test]
|
||||
fn test_integrate_num_vals_empty() {
|
||||
assert!(integrate_num_vals(iter::empty()).eq(iter::once(0)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_integrate_num_vals_one_el() {
|
||||
assert!(integrate_num_vals(iter::once(10)).eq([0, 10].into_iter()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_integrate_num_vals_several() {
|
||||
assert!(integrate_num_vals([3, 0, 10, 20].into_iter()).eq([0, 3, 3, 13, 33].into_iter()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_column_index_optional_shuffle() {
|
||||
let optional_index: ColumnIndex = OptionalIndex::for_test(2, &[0]).into();
|
||||
let column_indexes = vec![optional_index, ColumnIndex::Full];
|
||||
let row_addrs = vec![
|
||||
RowAddr {
|
||||
segment_ord: 0u32,
|
||||
row_id: 1u32,
|
||||
},
|
||||
RowAddr {
|
||||
segment_ord: 1u32,
|
||||
row_id: 0u32,
|
||||
},
|
||||
];
|
||||
let shuffle_merge_order = ShuffleMergeOrder::for_test(&[2, 1], row_addrs);
|
||||
let serializable_index = merge_column_index_shuffled(
|
||||
&column_indexes[..],
|
||||
Cardinality::Optional,
|
||||
&shuffle_merge_order,
|
||||
);
|
||||
let SerializableColumnIndex::Optional { non_null_row_ids, num_rows } = serializable_index else { panic!() };
|
||||
assert_eq!(num_rows, 2);
|
||||
let non_null_rows: Vec<RowId> = non_null_row_ids.boxed_iter().collect();
|
||||
assert_eq!(&non_null_rows, &[1]);
|
||||
}
|
||||
}
|
||||
@@ -5,18 +5,15 @@ mod serialize;
|
||||
|
||||
use std::ops::Range;
|
||||
|
||||
pub use merge::merge_column_index;
|
||||
pub use merge::stack_column_index;
|
||||
pub use optional_index::{OptionalIndex, Set};
|
||||
pub use serialize::{open_column_index, serialize_column_index, SerializableColumnIndex};
|
||||
|
||||
use crate::column_index::multivalued_index::MultiValueIndex;
|
||||
use crate::{Cardinality, DocId, RowId};
|
||||
use crate::{Cardinality, RowId};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
#[derive(Clone)]
|
||||
pub enum ColumnIndex {
|
||||
Empty {
|
||||
num_docs: u32,
|
||||
},
|
||||
Full,
|
||||
Optional(OptionalIndex),
|
||||
/// In addition, at index num_rows, an extra value is added
|
||||
@@ -24,125 +21,31 @@ pub enum ColumnIndex {
|
||||
Multivalued(MultiValueIndex),
|
||||
}
|
||||
|
||||
impl From<OptionalIndex> for ColumnIndex {
|
||||
fn from(optional_index: OptionalIndex) -> ColumnIndex {
|
||||
ColumnIndex::Optional(optional_index)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<MultiValueIndex> for ColumnIndex {
|
||||
fn from(multi_value_index: MultiValueIndex) -> ColumnIndex {
|
||||
ColumnIndex::Multivalued(multi_value_index)
|
||||
}
|
||||
}
|
||||
|
||||
impl ColumnIndex {
|
||||
// Returns the cardinality of the column index.
|
||||
//
|
||||
// By convention, if the column contains no docs, we consider that it is
|
||||
// full.
|
||||
#[inline]
|
||||
pub fn get_cardinality(&self) -> Cardinality {
|
||||
match self {
|
||||
ColumnIndex::Empty { num_docs: 0 } | ColumnIndex::Full => Cardinality::Full,
|
||||
ColumnIndex::Empty { .. } => Cardinality::Optional,
|
||||
ColumnIndex::Full => Cardinality::Full,
|
||||
ColumnIndex::Optional(_) => Cardinality::Optional,
|
||||
ColumnIndex::Multivalued(_) => Cardinality::Multivalued,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if and only if there are at least one value associated to the row.
|
||||
pub fn has_value(&self, doc_id: DocId) -> bool {
|
||||
pub fn value_row_ids(&self, row_id: RowId) -> Range<RowId> {
|
||||
match self {
|
||||
ColumnIndex::Empty { .. } => false,
|
||||
ColumnIndex::Full => true,
|
||||
ColumnIndex::Optional(optional_index) => optional_index.contains(doc_id),
|
||||
ColumnIndex::Multivalued(multivalued_index) => {
|
||||
!multivalued_index.range(doc_id).is_empty()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn value_row_ids(&self, doc_id: DocId) -> Range<RowId> {
|
||||
match self {
|
||||
ColumnIndex::Empty { .. } => 0..0,
|
||||
ColumnIndex::Full => doc_id..doc_id + 1,
|
||||
ColumnIndex::Full => row_id..row_id + 1,
|
||||
ColumnIndex::Optional(optional_index) => {
|
||||
if let Some(val) = optional_index.rank_if_exists(doc_id) {
|
||||
if let Some(val) = optional_index.rank_if_exists(row_id) {
|
||||
val..val + 1
|
||||
} else {
|
||||
0..0
|
||||
}
|
||||
}
|
||||
ColumnIndex::Multivalued(multivalued_index) => multivalued_index.range(doc_id),
|
||||
ColumnIndex::Multivalued(multivalued_index) => multivalued_index.range(row_id),
|
||||
}
|
||||
}
|
||||
|
||||
/// Translates a block of docis to row_ids.
|
||||
///
|
||||
/// returns the row_ids and the matching docids on the same index
|
||||
/// e.g.
|
||||
/// DocId In: [0, 5, 6]
|
||||
/// DocId Out: [0, 0, 6, 6]
|
||||
/// RowId Out: [0, 1, 2, 3]
|
||||
#[inline]
|
||||
pub fn docids_to_rowids(
|
||||
&self,
|
||||
doc_ids: &[DocId],
|
||||
doc_ids_out: &mut Vec<DocId>,
|
||||
row_ids: &mut Vec<RowId>,
|
||||
) {
|
||||
pub fn select_batch_in_place(&self, rank_ids: &mut Vec<RowId>) {
|
||||
match self {
|
||||
ColumnIndex::Empty { .. } => {}
|
||||
ColumnIndex::Full => {
|
||||
doc_ids_out.extend_from_slice(doc_ids);
|
||||
row_ids.extend_from_slice(doc_ids);
|
||||
}
|
||||
ColumnIndex::Optional(optional_index) => {
|
||||
for doc_id in doc_ids {
|
||||
if let Some(row_id) = optional_index.rank_if_exists(*doc_id) {
|
||||
doc_ids_out.push(*doc_id);
|
||||
row_ids.push(row_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
ColumnIndex::Multivalued(multivalued_index) => {
|
||||
for doc_id in doc_ids {
|
||||
for row_id in multivalued_index.range(*doc_id) {
|
||||
doc_ids_out.push(*doc_id);
|
||||
row_ids.push(row_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn docid_range_to_rowids(&self, doc_id: Range<DocId>) -> Range<RowId> {
|
||||
match self {
|
||||
ColumnIndex::Empty { .. } => 0..0,
|
||||
ColumnIndex::Full => doc_id,
|
||||
ColumnIndex::Optional(optional_index) => {
|
||||
let row_start = optional_index.rank(doc_id.start);
|
||||
let row_end = optional_index.rank(doc_id.end);
|
||||
row_start..row_end
|
||||
}
|
||||
ColumnIndex::Multivalued(multivalued_index) => {
|
||||
let end_docid = doc_id.end.min(multivalued_index.num_docs() - 1) + 1;
|
||||
let start_docid = doc_id.start.min(end_docid);
|
||||
|
||||
let row_start = multivalued_index.start_index_column.get_val(start_docid);
|
||||
let row_end = multivalued_index.start_index_column.get_val(end_docid);
|
||||
|
||||
row_start..row_end
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn select_batch_in_place(&self, doc_id_start: DocId, rank_ids: &mut Vec<RowId>) {
|
||||
match self {
|
||||
ColumnIndex::Empty { .. } => {
|
||||
rank_ids.clear();
|
||||
}
|
||||
ColumnIndex::Full => {
|
||||
// No need to do anything:
|
||||
// value_idx and row_idx are the same.
|
||||
@@ -151,26 +54,9 @@ impl ColumnIndex {
|
||||
optional_index.select_batch(&mut rank_ids[..]);
|
||||
}
|
||||
ColumnIndex::Multivalued(multivalued_index) => {
|
||||
multivalued_index.select_batch_in_place(doc_id_start, rank_ids)
|
||||
// TODO important: avoid using 0u32, and restart from the beginning all of the time.
|
||||
multivalued_index.select_batch_in_place(0u32, rank_ids)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::{Cardinality, ColumnIndex};
|
||||
|
||||
#[test]
|
||||
fn test_column_index_get_cardinality() {
|
||||
assert_eq!(
|
||||
ColumnIndex::Empty { num_docs: 0 }.get_cardinality(),
|
||||
Cardinality::Full
|
||||
);
|
||||
assert_eq!(ColumnIndex::Full.get_cardinality(), Cardinality::Full);
|
||||
assert_eq!(
|
||||
ColumnIndex::Empty { num_docs: 1 }.get_cardinality(),
|
||||
Cardinality::Optional
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,17 +5,16 @@ use std::sync::Arc;
|
||||
|
||||
use common::OwnedBytes;
|
||||
|
||||
use crate::column_values::{
|
||||
load_u64_based_column_values, serialize_u64_based_column_values, CodecType, ColumnValues,
|
||||
};
|
||||
use crate::column_values::u64_based::CodecType;
|
||||
use crate::column_values::ColumnValues;
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{DocId, RowId};
|
||||
use crate::RowId;
|
||||
|
||||
pub fn serialize_multivalued_index(
|
||||
multivalued_index: &dyn Iterable<RowId>,
|
||||
output: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
serialize_u64_based_column_values(
|
||||
crate::column_values::u64_based::serialize_u64_based_column_values(
|
||||
multivalued_index,
|
||||
&[CodecType::Bitpacked, CodecType::Linear],
|
||||
output,
|
||||
@@ -24,7 +23,8 @@ pub fn serialize_multivalued_index(
|
||||
}
|
||||
|
||||
pub fn open_multivalued_index(bytes: OwnedBytes) -> io::Result<MultiValueIndex> {
|
||||
let start_index_column: Arc<dyn ColumnValues<RowId>> = load_u64_based_column_values(bytes)?;
|
||||
let start_index_column: Arc<dyn ColumnValues<RowId>> =
|
||||
crate::column_values::u64_based::load_u64_based_column_values(bytes)?;
|
||||
Ok(MultiValueIndex { start_index_column })
|
||||
}
|
||||
|
||||
@@ -35,14 +35,6 @@ pub struct MultiValueIndex {
|
||||
pub start_index_column: Arc<dyn crate::ColumnValues<RowId>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for MultiValueIndex {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
f.debug_struct("MultiValuedIndex")
|
||||
.field("num_rows", &self.start_index_column.num_vals())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Arc<dyn ColumnValues<RowId>>> for MultiValueIndex {
|
||||
fn from(start_index_column: Arc<dyn ColumnValues<RowId>>) -> Self {
|
||||
MultiValueIndex { start_index_column }
|
||||
@@ -50,30 +42,23 @@ impl From<Arc<dyn ColumnValues<RowId>>> for MultiValueIndex {
|
||||
}
|
||||
|
||||
impl MultiValueIndex {
|
||||
pub fn for_test(start_offsets: &[RowId]) -> MultiValueIndex {
|
||||
let mut buffer = Vec::new();
|
||||
serialize_multivalued_index(&start_offsets, &mut buffer).unwrap();
|
||||
let bytes = OwnedBytes::new(buffer);
|
||||
open_multivalued_index(bytes).unwrap()
|
||||
}
|
||||
|
||||
/// Returns `[start, end)`, such that the values associated with
|
||||
/// the given document are `start..end`.
|
||||
#[inline]
|
||||
pub(crate) fn range(&self, doc_id: DocId) -> Range<RowId> {
|
||||
let start = self.start_index_column.get_val(doc_id);
|
||||
let end = self.start_index_column.get_val(doc_id + 1);
|
||||
pub(crate) fn range(&self, row_id: RowId) -> Range<RowId> {
|
||||
let start = self.start_index_column.get_val(row_id);
|
||||
let end = self.start_index_column.get_val(row_id + 1);
|
||||
start..end
|
||||
}
|
||||
|
||||
/// Returns the number of documents in the index.
|
||||
#[inline]
|
||||
pub fn num_docs(&self) -> u32 {
|
||||
pub fn num_rows(&self) -> u32 {
|
||||
self.start_index_column.num_vals() - 1
|
||||
}
|
||||
|
||||
/// Converts a list of ranks (row ids of values) in a 1:n index to the corresponding list of
|
||||
/// docids. Positions are converted inplace to docids.
|
||||
/// row_ids. Positions are converted inplace to docids.
|
||||
///
|
||||
/// Since there is no index for value pos -> docid, but docid -> value pos range, we scan the
|
||||
/// index.
|
||||
@@ -84,20 +69,20 @@ impl MultiValueIndex {
|
||||
/// TODO: Instead of a linear scan we can employ a exponential search into binary search to
|
||||
/// match a docid to its value position.
|
||||
#[allow(clippy::bool_to_int_with_if)]
|
||||
pub(crate) fn select_batch_in_place(&self, docid_start: DocId, ranks: &mut Vec<u32>) {
|
||||
pub(crate) fn select_batch_in_place(&self, row_start: RowId, ranks: &mut Vec<u32>) {
|
||||
if ranks.is_empty() {
|
||||
return;
|
||||
}
|
||||
let mut cur_doc = docid_start;
|
||||
let mut cur_doc = row_start;
|
||||
let mut last_doc = None;
|
||||
|
||||
assert!(self.start_index_column.get_val(docid_start) <= ranks[0]);
|
||||
assert!(self.start_index_column.get_val(row_start) as u32 <= ranks[0]);
|
||||
|
||||
let mut write_doc_pos = 0;
|
||||
for i in 0..ranks.len() {
|
||||
let pos = ranks[i];
|
||||
loop {
|
||||
let end = self.start_index_column.get_val(cur_doc + 1);
|
||||
let end = self.start_index_column.get_val(cur_doc + 1) as u32;
|
||||
if end > pos {
|
||||
ranks[write_doc_pos] = cur_doc;
|
||||
write_doc_pos += if last_doc == Some(cur_doc) { 0 } else { 1 };
|
||||
@@ -114,8 +99,11 @@ impl MultiValueIndex {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::ops::Range;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::MultiValueIndex;
|
||||
use crate::column_values::IterColumn;
|
||||
use crate::{ColumnValues, RowId};
|
||||
|
||||
fn index_to_pos_helper(
|
||||
index: &MultiValueIndex,
|
||||
@@ -129,8 +117,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_positions_to_docid() {
|
||||
let index = MultiValueIndex::for_test(&[0, 10, 12, 15, 22, 23]);
|
||||
assert_eq!(index.num_docs(), 5);
|
||||
let offsets: Vec<RowId> = vec![0, 10, 12, 15, 22, 23]; // docid values are [0..10, 10..12, 12..15, etc.]
|
||||
let column: Arc<dyn ColumnValues<RowId>> = Arc::new(IterColumn::from(offsets.into_iter()));
|
||||
let index = MultiValueIndex::from(column);
|
||||
assert_eq!(index.num_rows(), 5);
|
||||
let positions = &[10u32, 11, 15, 20, 21, 22];
|
||||
assert_eq!(index_to_pos_helper(&index, 0..5, positions), vec![1, 3, 4]);
|
||||
assert_eq!(index_to_pos_helper(&index, 1..5, positions), vec![1, 3, 4]);
|
||||
|
||||
@@ -11,7 +11,7 @@ use set_block::{
|
||||
};
|
||||
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{DocId, InvalidData, RowId};
|
||||
use crate::{InvalidData, RowId};
|
||||
|
||||
/// The threshold for for number of elements after which we switch to dense block encoding.
|
||||
///
|
||||
@@ -88,12 +88,19 @@ pub struct OptionalIndex {
|
||||
block_metas: Arc<[BlockMeta]>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for OptionalIndex {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("OptionalIndex")
|
||||
.field("num_rows", &self.num_rows)
|
||||
.field("num_non_null_rows", &self.num_non_null_rows)
|
||||
.finish_non_exhaustive()
|
||||
impl OptionalIndex {
|
||||
pub fn num_rows(&self) -> RowId {
|
||||
self.num_rows
|
||||
}
|
||||
|
||||
pub fn num_non_nulls(&self) -> RowId {
|
||||
self.num_non_null_rows
|
||||
}
|
||||
|
||||
pub fn iter_rows<'a>(&'a self) -> impl Iterator<Item = RowId> + 'a {
|
||||
// TODO optimize
|
||||
let mut select_batch = self.select_cursor();
|
||||
(0..self.num_non_null_rows).map(move |rank| select_batch.select(rank))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -186,11 +193,11 @@ impl Set<RowId> for OptionalIndex {
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn rank(&self, doc_id: DocId) -> RowId {
|
||||
fn rank(&self, row_id: RowId) -> RowId {
|
||||
let RowAddr {
|
||||
block_id,
|
||||
in_block_row_id,
|
||||
} = row_addr_from_row_id(doc_id);
|
||||
} = row_addr_from_row_id(row_id);
|
||||
let block_meta = self.block_metas[block_id as usize];
|
||||
let block = self.block(block_meta);
|
||||
let block_offset_row_id = match block {
|
||||
@@ -201,11 +208,11 @@ impl Set<RowId> for OptionalIndex {
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn rank_if_exists(&self, doc_id: DocId) -> Option<RowId> {
|
||||
fn rank_if_exists(&self, row_id: RowId) -> Option<RowId> {
|
||||
let RowAddr {
|
||||
block_id,
|
||||
in_block_row_id,
|
||||
} = row_addr_from_row_id(doc_id);
|
||||
} = row_addr_from_row_id(row_id);
|
||||
let block_meta = self.block_metas[block_id as usize];
|
||||
let block = self.block(block_meta);
|
||||
let block_offset_row_id = match block {
|
||||
@@ -229,7 +236,7 @@ impl Set<RowId> for OptionalIndex {
|
||||
block_doc_idx_start + in_block_rank as u32
|
||||
}
|
||||
|
||||
fn select_cursor(&self) -> OptionalIndexSelectCursor<'_> {
|
||||
fn select_cursor<'b>(&'b self) -> OptionalIndexSelectCursor<'b> {
|
||||
OptionalIndexSelectCursor {
|
||||
current_block_cursor: BlockSelectCursor::Sparse(
|
||||
SparseBlockCodec::open(b"").select_cursor(),
|
||||
@@ -244,31 +251,6 @@ impl Set<RowId> for OptionalIndex {
|
||||
}
|
||||
|
||||
impl OptionalIndex {
|
||||
pub fn for_test(num_rows: RowId, row_ids: &[RowId]) -> OptionalIndex {
|
||||
assert!(row_ids
|
||||
.last()
|
||||
.copied()
|
||||
.map(|last_row_id| last_row_id < num_rows)
|
||||
.unwrap_or(true));
|
||||
let mut buffer = Vec::new();
|
||||
serialize_optional_index(&row_ids, num_rows, &mut buffer).unwrap();
|
||||
let bytes = OwnedBytes::new(buffer);
|
||||
open_optional_index(bytes).unwrap()
|
||||
}
|
||||
|
||||
pub fn num_docs(&self) -> RowId {
|
||||
self.num_rows
|
||||
}
|
||||
|
||||
pub fn num_non_nulls(&self) -> RowId {
|
||||
self.num_non_null_rows
|
||||
}
|
||||
|
||||
pub fn iter_rows(&self) -> impl Iterator<Item = RowId> + '_ {
|
||||
// TODO optimize
|
||||
let mut select_batch = self.select_cursor();
|
||||
(0..self.num_non_null_rows).map(move |rank| select_batch.select(rank))
|
||||
}
|
||||
pub fn select_batch(&self, ranks: &mut [RowId]) {
|
||||
let mut select_cursor = self.select_cursor();
|
||||
for rank in ranks.iter_mut() {
|
||||
@@ -277,7 +259,7 @@ impl OptionalIndex {
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn block(&self, block_meta: BlockMeta) -> Block<'_> {
|
||||
fn block<'a>(&'a self, block_meta: BlockMeta) -> Block<'a> {
|
||||
let BlockMeta {
|
||||
start_byte_offset,
|
||||
block_variant,
|
||||
@@ -360,7 +342,7 @@ fn serialize_optional_index_block(block_els: &[u16], out: &mut impl io::Write) -
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn serialize_optional_index<W: io::Write>(
|
||||
pub fn serialize_optional_index<'a, W: io::Write>(
|
||||
non_null_rows: &dyn Iterable<RowId>,
|
||||
num_rows: RowId,
|
||||
output: &mut W,
|
||||
@@ -436,7 +418,7 @@ impl SerializedBlockMeta {
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn to_bytes(self) -> [u8; SERIALIZED_BLOCK_META_NUM_BYTES] {
|
||||
fn to_bytes(&self) -> [u8; SERIALIZED_BLOCK_META_NUM_BYTES] {
|
||||
assert!(self.num_non_null_rows > 0);
|
||||
let mut bytes = [0u8; SERIALIZED_BLOCK_META_NUM_BYTES];
|
||||
bytes[0..2].copy_from_slice(&self.block_id.to_le_bytes());
|
||||
@@ -449,7 +431,7 @@ impl SerializedBlockMeta {
|
||||
|
||||
#[inline]
|
||||
fn is_sparse(num_rows_in_block: u32) -> bool {
|
||||
num_rows_in_block < DENSE_BLOCK_THRESHOLD
|
||||
num_rows_in_block < DENSE_BLOCK_THRESHOLD as u32
|
||||
}
|
||||
|
||||
fn deserialize_optional_index_block_metadatas(
|
||||
@@ -457,7 +439,7 @@ fn deserialize_optional_index_block_metadatas(
|
||||
num_rows: u32,
|
||||
) -> (Box<[BlockMeta]>, u32) {
|
||||
let num_blocks = data.len() / SERIALIZED_BLOCK_META_NUM_BYTES;
|
||||
let mut block_metas = Vec::with_capacity(num_blocks + 1);
|
||||
let mut block_metas = Vec::with_capacity(num_blocks as usize + 1);
|
||||
let mut start_byte_offset = 0;
|
||||
let mut non_null_rows_before_block = 0;
|
||||
for block_meta_bytes in data.chunks_exact(SERIALIZED_BLOCK_META_NUM_BYTES) {
|
||||
@@ -488,7 +470,7 @@ fn deserialize_optional_index_block_metadatas(
|
||||
block_variant,
|
||||
});
|
||||
start_byte_offset += block_variant.num_bytes_in_block();
|
||||
non_null_rows_before_block += num_non_null_rows;
|
||||
non_null_rows_before_block += num_non_null_rows as u32;
|
||||
}
|
||||
block_metas.resize(
|
||||
((num_rows + BLOCK_SIZE - 1) / BLOCK_SIZE) as usize,
|
||||
@@ -510,7 +492,7 @@ pub fn open_optional_index(bytes: OwnedBytes) -> io::Result<OptionalIndex> {
|
||||
num_non_empty_block_bytes as usize * SERIALIZED_BLOCK_META_NUM_BYTES;
|
||||
let (block_data, block_metas) = bytes.rsplit(block_metas_num_bytes);
|
||||
let (block_metas, num_non_null_rows) =
|
||||
deserialize_optional_index_block_metadatas(block_metas.as_slice(), num_rows);
|
||||
deserialize_optional_index_block_metadatas(block_metas.as_slice(), num_rows).into();
|
||||
let optional_index = OptionalIndex {
|
||||
num_rows,
|
||||
num_non_null_rows,
|
||||
|
||||
@@ -10,7 +10,7 @@ pub trait SetCodec {
|
||||
///
|
||||
/// May panic if the elements are not sorted.
|
||||
fn serialize(els: impl Iterator<Item = Self::Item>, wrt: impl io::Write) -> io::Result<()>;
|
||||
fn open(data: &[u8]) -> Self::Reader<'_>;
|
||||
fn open<'a>(data: &'a [u8]) -> Self::Reader<'a>;
|
||||
}
|
||||
|
||||
/// Stateful object that makes it possible to compute several select in a row,
|
||||
@@ -43,5 +43,5 @@ pub trait Set<T> {
|
||||
fn select(&self, rank: T) -> T;
|
||||
|
||||
/// Creates a brand new select cursor.
|
||||
fn select_cursor(&self) -> Self::SelectCursor<'_>;
|
||||
fn select_cursor<'b>(&'b self) -> Self::SelectCursor<'b>;
|
||||
}
|
||||
|
||||
@@ -32,7 +32,7 @@ pub const MINI_BLOCK_NUM_BYTES: usize = MINI_BLOCK_BITVEC_NUM_BYTES + MINI_BLOCK
|
||||
|
||||
/// Number of bytes in a dense block.
|
||||
pub const DENSE_BLOCK_NUM_BYTES: u32 =
|
||||
(ELEMENTS_PER_BLOCK / ELEMENTS_PER_MINI_BLOCK as u32) * MINI_BLOCK_NUM_BYTES as u32;
|
||||
(ELEMENTS_PER_BLOCK as u32 / ELEMENTS_PER_MINI_BLOCK as u32) * MINI_BLOCK_NUM_BYTES as u32;
|
||||
|
||||
pub struct DenseBlockCodec;
|
||||
|
||||
@@ -45,7 +45,7 @@ impl SetCodec for DenseBlockCodec {
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn open(data: &[u8]) -> Self::Reader<'_> {
|
||||
fn open<'a>(data: &'a [u8]) -> Self::Reader<'a> {
|
||||
assert_eq!(data.len(), DENSE_BLOCK_NUM_BYTES as usize);
|
||||
DenseBlock(data)
|
||||
}
|
||||
@@ -94,7 +94,7 @@ impl DenseMiniBlock {
|
||||
Self { bitvec, rank }
|
||||
}
|
||||
|
||||
fn to_bytes(self) -> [u8; MINI_BLOCK_NUM_BYTES] {
|
||||
fn to_bytes(&self) -> [u8; MINI_BLOCK_NUM_BYTES] {
|
||||
let mut bytes = [0u8; MINI_BLOCK_NUM_BYTES];
|
||||
bytes[..MINI_BLOCK_BITVEC_NUM_BYTES].copy_from_slice(&self.bitvec.to_le_bytes());
|
||||
bytes[MINI_BLOCK_BITVEC_NUM_BYTES..].copy_from_slice(&self.rank.to_le_bytes());
|
||||
@@ -166,7 +166,7 @@ impl<'a> Set<u16> for DenseBlock<'a> {
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn select_cursor(&self) -> Self::SelectCursor<'_> {
|
||||
fn select_cursor<'b>(&'b self) -> Self::SelectCursor<'b> {
|
||||
DenseBlockSelectCursor {
|
||||
block_id: 0,
|
||||
dense_block: *self,
|
||||
@@ -229,7 +229,7 @@ pub fn serialize_dense_codec(
|
||||
while block_id > current_block_id {
|
||||
let dense_mini_block = DenseMiniBlock {
|
||||
bitvec: block,
|
||||
rank: non_null_rows_before,
|
||||
rank: non_null_rows_before as u16,
|
||||
};
|
||||
output.write_all(&dense_mini_block.to_bytes())?;
|
||||
non_null_rows_before += block.count_ones() as u16;
|
||||
|
||||
@@ -16,7 +16,7 @@ impl SetCodec for SparseBlockCodec {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn open(data: &[u8]) -> Self::Reader<'_> {
|
||||
fn open<'a>(data: &'a [u8]) -> Self::Reader<'a> {
|
||||
SparseBlock(data)
|
||||
}
|
||||
}
|
||||
@@ -56,7 +56,7 @@ impl<'a> Set<u16> for SparseBlock<'a> {
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn select_cursor(&self) -> Self::SelectCursor<'_> {
|
||||
fn select_cursor<'b>(&'b self) -> Self::SelectCursor<'b> {
|
||||
*self
|
||||
}
|
||||
}
|
||||
|
||||
@@ -37,7 +37,7 @@ proptest! {
|
||||
fn test_with_random_sets_simple() {
|
||||
let vals = 10..BLOCK_SIZE * 2;
|
||||
let mut out: Vec<u8> = Vec::new();
|
||||
serialize_optional_index(&vals, 100, &mut out).unwrap();
|
||||
serialize_optional_index(&vals.clone(), 100, &mut out).unwrap();
|
||||
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
|
||||
let ranks: Vec<u32> = (65_472u32..65_473u32).collect();
|
||||
let els: Vec<u32> = ranks.iter().copied().map(|rank| rank + 10).collect();
|
||||
@@ -107,43 +107,59 @@ fn test_null_index(data: &[bool]) {
|
||||
|
||||
#[test]
|
||||
fn test_optional_index_test_translation() {
|
||||
let optional_index = OptionalIndex::for_test(4, &[0, 2]);
|
||||
let mut select_cursor = optional_index.select_cursor();
|
||||
let mut out = vec![];
|
||||
let iter = &[true, false, true, false];
|
||||
serialize_optional_index(&&iter[..], iter.len() as u32, &mut out).unwrap();
|
||||
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
|
||||
let mut select_cursor = null_index.select_cursor();
|
||||
assert_eq!(select_cursor.select(0), 0);
|
||||
assert_eq!(select_cursor.select(1), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optional_index_translate() {
|
||||
let optional_index = OptionalIndex::for_test(4, &[0, 2]);
|
||||
assert_eq!(optional_index.rank_if_exists(0), Some(0));
|
||||
assert_eq!(optional_index.rank_if_exists(2), Some(1));
|
||||
let mut out = vec![];
|
||||
let iter = &[true, false, true, false];
|
||||
serialize_optional_index(&&iter[..], iter.len() as RowId, &mut out).unwrap();
|
||||
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
|
||||
assert_eq!(null_index.rank_if_exists(0), Some(0));
|
||||
assert_eq!(null_index.rank_if_exists(2), Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optional_index_small() {
|
||||
let optional_index = OptionalIndex::for_test(4, &[0, 2]);
|
||||
assert!(optional_index.contains(0));
|
||||
assert!(!optional_index.contains(1));
|
||||
assert!(optional_index.contains(2));
|
||||
assert!(!optional_index.contains(3));
|
||||
let mut out = vec![];
|
||||
let iter = &[true, false, true, false];
|
||||
serialize_optional_index(&&iter[..], iter.len() as RowId, &mut out).unwrap();
|
||||
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
|
||||
assert!(null_index.contains(0));
|
||||
assert!(!null_index.contains(1));
|
||||
assert!(null_index.contains(2));
|
||||
assert!(!null_index.contains(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optional_index_large() {
|
||||
let row_ids = &[ELEMENTS_PER_BLOCK, ELEMENTS_PER_BLOCK + 1];
|
||||
let optional_index = OptionalIndex::for_test(ELEMENTS_PER_BLOCK + 2, row_ids);
|
||||
assert!(!optional_index.contains(0));
|
||||
assert!(!optional_index.contains(100));
|
||||
assert!(!optional_index.contains(ELEMENTS_PER_BLOCK - 1));
|
||||
assert!(optional_index.contains(ELEMENTS_PER_BLOCK));
|
||||
assert!(optional_index.contains(ELEMENTS_PER_BLOCK + 1));
|
||||
let mut docs = vec![];
|
||||
docs.extend((0..ELEMENTS_PER_BLOCK).map(|_idx| false));
|
||||
docs.extend((0..=1).map(|_idx| true));
|
||||
|
||||
let mut out = vec![];
|
||||
serialize_optional_index(&&docs[..], docs.len() as RowId, &mut out).unwrap();
|
||||
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
|
||||
assert!(!null_index.contains(0));
|
||||
assert!(!null_index.contains(100));
|
||||
assert!(!null_index.contains(ELEMENTS_PER_BLOCK - 1));
|
||||
assert!(null_index.contains(ELEMENTS_PER_BLOCK));
|
||||
assert!(null_index.contains(ELEMENTS_PER_BLOCK + 1));
|
||||
}
|
||||
|
||||
fn test_optional_index_iter_aux(row_ids: &[RowId], num_rows: RowId) {
|
||||
let optional_index = OptionalIndex::for_test(num_rows, row_ids);
|
||||
assert_eq!(optional_index.num_docs(), num_rows);
|
||||
assert!(optional_index.iter_rows().eq(row_ids.iter().copied()));
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
serialize_optional_index(&row_ids, num_rows, &mut buffer).unwrap();
|
||||
let null_index = open_optional_index(OwnedBytes::new(buffer)).unwrap();
|
||||
assert_eq!(null_index.num_rows(), num_rows);
|
||||
assert!(null_index.iter_rows().eq(row_ids.iter().copied()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -152,9 +168,11 @@ fn test_optional_index_iter_empty() {
|
||||
}
|
||||
|
||||
fn test_optional_index_rank_aux(row_ids: &[RowId]) {
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
let num_rows = row_ids.last().copied().unwrap_or(0u32) + 1;
|
||||
let null_index = OptionalIndex::for_test(num_rows, row_ids);
|
||||
assert_eq!(null_index.num_docs(), num_rows);
|
||||
serialize_optional_index(&row_ids, num_rows, &mut buffer).unwrap();
|
||||
let null_index = open_optional_index(OwnedBytes::new(buffer)).unwrap();
|
||||
assert_eq!(null_index.num_rows(), num_rows);
|
||||
for (row_id, row_val) in row_ids.iter().copied().enumerate() {
|
||||
assert_eq!(null_index.rank(row_val), row_id as u32);
|
||||
assert_eq!(null_index.rank_if_exists(row_val), Some(row_id as u32));
|
||||
@@ -189,16 +207,6 @@ fn test_optional_index_iter_dense_block() {
|
||||
test_optional_index_iter_aux(&block, 3 * BLOCK_SIZE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optional_index_for_tests() {
|
||||
let optional_index = OptionalIndex::for_test(4, &[1, 2]);
|
||||
assert!(!optional_index.contains(0));
|
||||
assert!(optional_index.contains(1));
|
||||
assert!(optional_index.contains(2));
|
||||
assert!(!optional_index.contains(3));
|
||||
assert_eq!(optional_index.num_docs(), 4);
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "unstable"))]
|
||||
mod bench {
|
||||
|
||||
@@ -212,13 +220,10 @@ mod bench {
|
||||
fn gen_bools(fill_ratio: f64) -> OptionalIndex {
|
||||
let mut out = Vec::new();
|
||||
let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
|
||||
let vals: Vec<RowId> = (0..TOTAL_NUM_VALUES)
|
||||
let vals: Vec<bool> = (0..TOTAL_NUM_VALUES)
|
||||
.map(|_| rng.gen_bool(fill_ratio))
|
||||
.enumerate()
|
||||
.filter(|(pos, val)| *val)
|
||||
.map(|(pos, _)| pos as RowId)
|
||||
.collect();
|
||||
serialize_optional_index(&&vals[..], TOTAL_NUM_VALUES, &mut out).unwrap();
|
||||
serialize_optional_index(&&vals[..], &mut out).unwrap();
|
||||
let codec = open_optional_index(OwnedBytes::new(out)).unwrap();
|
||||
codec
|
||||
}
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::OwnedBytes;
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use test::{self, Bencher};
|
||||
|
||||
use super::*;
|
||||
use crate::column_values::u64_based::*;
|
||||
|
||||
fn get_data() -> Vec<u64> {
|
||||
let mut rng = StdRng::seed_from_u64(2u64);
|
||||
let mut data: Vec<_> = (100..55000_u64)
|
||||
.map(|num| num + rng.gen::<u8>() as u64)
|
||||
.collect();
|
||||
data.push(99_000);
|
||||
data.insert(1000, 2000);
|
||||
data.insert(2000, 100);
|
||||
data.insert(3000, 4100);
|
||||
data.insert(4000, 100);
|
||||
data.insert(5000, 800);
|
||||
data
|
||||
}
|
||||
|
||||
fn compute_stats(vals: impl Iterator<Item = u64>) -> ColumnStats {
|
||||
let mut stats_collector = StatsCollector::default();
|
||||
for val in vals {
|
||||
stats_collector.collect(val);
|
||||
}
|
||||
stats_collector.stats()
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn value_iter() -> impl Iterator<Item = u64> {
|
||||
0..20_000
|
||||
}
|
||||
fn get_reader_for_bench<Codec: ColumnCodec>(data: &[u64]) -> Codec::ColumnValues {
|
||||
let mut bytes = Vec::new();
|
||||
let stats = compute_stats(data.iter().cloned());
|
||||
let mut codec_serializer = Codec::estimator();
|
||||
for val in data {
|
||||
codec_serializer.collect(*val);
|
||||
}
|
||||
codec_serializer.serialize(&stats, Box::new(data.iter().copied()).as_mut(), &mut bytes);
|
||||
|
||||
Codec::load(OwnedBytes::new(bytes)).unwrap()
|
||||
}
|
||||
fn bench_get<Codec: ColumnCodec>(b: &mut Bencher, data: &[u64]) {
|
||||
let col = get_reader_for_bench::<Codec>(data);
|
||||
b.iter(|| {
|
||||
let mut sum = 0u64;
|
||||
for pos in value_iter() {
|
||||
let val = col.get_val(pos as u32);
|
||||
sum = sum.wrapping_add(val);
|
||||
}
|
||||
sum
|
||||
});
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn bench_get_dynamic_helper(b: &mut Bencher, col: Arc<dyn ColumnValues>) {
|
||||
b.iter(|| {
|
||||
let mut sum = 0u64;
|
||||
for pos in value_iter() {
|
||||
let val = col.get_val(pos as u32);
|
||||
sum = sum.wrapping_add(val);
|
||||
}
|
||||
sum
|
||||
});
|
||||
}
|
||||
|
||||
fn bench_get_dynamic<Codec: ColumnCodec>(b: &mut Bencher, data: &[u64]) {
|
||||
let col = Arc::new(get_reader_for_bench::<Codec>(data));
|
||||
bench_get_dynamic_helper(b, col);
|
||||
}
|
||||
fn bench_create<Codec: ColumnCodec>(b: &mut Bencher, data: &[u64]) {
|
||||
let stats = compute_stats(data.iter().cloned());
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
b.iter(|| {
|
||||
bytes.clear();
|
||||
let mut codec_serializer = Codec::estimator();
|
||||
for val in data.iter().take(1024) {
|
||||
codec_serializer.collect(*val);
|
||||
}
|
||||
|
||||
codec_serializer.serialize(&stats, Box::new(data.iter().copied()).as_mut(), &mut bytes)
|
||||
});
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_fastfield_bitpack_create(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_create::<BitpackedCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_linearinterpol_create(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_create::<LinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_multilinearinterpol_create(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_create::<BlockwiseLinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_bitpack_get(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get::<BitpackedCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_bitpack_get_dynamic(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get_dynamic::<BitpackedCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_linearinterpol_get(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get::<LinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_linearinterpol_get_dynamic(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get_dynamic::<LinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_multilinearinterpol_get(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get::<BlockwiseLinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_multilinearinterpol_get_dynamic(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get_dynamic::<BlockwiseLinearCodec>(b, &data);
|
||||
}
|
||||
390
columnar/src/column_values/column.rs
Normal file
390
columnar/src/column_values/column.rs
Normal file
@@ -0,0 +1,390 @@
|
||||
use std::fmt::Debug;
|
||||
use std::marker::PhantomData;
|
||||
use std::ops::{Range, RangeInclusive};
|
||||
use std::sync::Arc;
|
||||
|
||||
use tantivy_bitpacker::minmax;
|
||||
|
||||
use crate::column_values::monotonic_mapping::StrictlyMonotonicFn;
|
||||
use crate::iterable::Iterable;
|
||||
|
||||
/// `ColumnValues` provides access to a dense field column.
|
||||
///
|
||||
/// `Column` are just a wrapper over `ColumnValues` and a `ColumnIndex`.
|
||||
pub trait ColumnValues<T: PartialOrd = u64>: Send + Sync {
|
||||
/// Return the value associated with the given idx.
|
||||
///
|
||||
/// This accessor should return as fast as possible.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// May panic if `idx` is greater than the column length.
|
||||
fn get_val(&self, idx: u32) -> T;
|
||||
|
||||
/// Fills an output buffer with the fast field values
|
||||
/// associated with the `DocId` going from
|
||||
/// `start` to `start + output.len()`.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Must panic if `start + output.len()` is greater than
|
||||
/// the segment's `maxdoc`.
|
||||
#[inline(always)]
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
for (out, idx) in output.iter_mut().zip(start..) {
|
||||
*out = self.get_val(idx as u32);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the positions of values which are in the provided value range.
|
||||
///
|
||||
/// Note that position == docid for single value fast fields
|
||||
#[inline(always)]
|
||||
fn get_docids_for_value_range(
|
||||
&self,
|
||||
value_range: RangeInclusive<T>,
|
||||
doc_id_range: Range<u32>,
|
||||
positions: &mut Vec<u32>,
|
||||
) {
|
||||
let doc_id_range = doc_id_range.start..doc_id_range.end.min(self.num_vals());
|
||||
for idx in doc_id_range.start..doc_id_range.end {
|
||||
let val = self.get_val(idx);
|
||||
if value_range.contains(&val) {
|
||||
positions.push(idx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the minimum value for this fast field.
|
||||
///
|
||||
/// This min_value may not be exact.
|
||||
/// For instance, the min value does not take in account of possible
|
||||
/// deleted document. All values are however guaranteed to be higher than
|
||||
/// `.min_value()`.
|
||||
fn min_value(&self) -> T;
|
||||
|
||||
/// Returns the maximum value for this fast field.
|
||||
///
|
||||
/// This max_value may not be exact.
|
||||
/// For instance, the max value does not take in account of possible
|
||||
/// deleted document. All values are however guaranteed to be higher than
|
||||
/// `.max_value()`.
|
||||
fn max_value(&self) -> T;
|
||||
|
||||
/// The number of values in the column.
|
||||
fn num_vals(&self) -> u32;
|
||||
|
||||
/// Returns a iterator over the data
|
||||
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = T> + 'a> {
|
||||
Box::new((0..self.num_vals()).map(|idx| self.get_val(idx)))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, T: PartialOrd> Iterable<T> for &'a [Arc<dyn ColumnValues<T>>] {
|
||||
fn boxed_iter(&self) -> Box<dyn Iterator<Item = T> + '_> {
|
||||
Box::new(self.iter().flat_map(|column_value| column_value.iter()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Copy + PartialOrd + Debug> ColumnValues<T> for Arc<dyn ColumnValues<T>> {
|
||||
#[inline(always)]
|
||||
fn get_val(&self, idx: u32) -> T {
|
||||
self.as_ref().get_val(idx)
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn min_value(&self) -> T {
|
||||
self.as_ref().min_value()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn max_value(&self) -> T {
|
||||
self.as_ref().max_value()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.as_ref().num_vals()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn iter<'b>(&'b self) -> Box<dyn Iterator<Item = T> + 'b> {
|
||||
self.as_ref().iter()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
self.as_ref().get_range(start, output)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, C: ColumnValues<T> + ?Sized, T: Copy + PartialOrd + Debug> ColumnValues<T> for &'a C {
|
||||
fn get_val(&self, idx: u32) -> T {
|
||||
(*self).get_val(idx)
|
||||
}
|
||||
|
||||
fn min_value(&self) -> T {
|
||||
(*self).min_value()
|
||||
}
|
||||
|
||||
fn max_value(&self) -> T {
|
||||
(*self).max_value()
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
(*self).num_vals()
|
||||
}
|
||||
|
||||
fn iter<'b>(&'b self) -> Box<dyn Iterator<Item = T> + 'b> {
|
||||
(*self).iter()
|
||||
}
|
||||
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
(*self).get_range(start, output)
|
||||
}
|
||||
}
|
||||
|
||||
/// VecColumn provides `Column` over a slice.
|
||||
pub struct VecColumn<'a, T = u64> {
|
||||
pub(crate) values: &'a [T],
|
||||
pub(crate) min_value: T,
|
||||
pub(crate) max_value: T,
|
||||
}
|
||||
|
||||
impl<'a, T: Copy + PartialOrd + Send + Sync + Debug> ColumnValues<T> for VecColumn<'a, T> {
|
||||
fn get_val(&self, position: u32) -> T {
|
||||
self.values[position as usize]
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = T> + '_> {
|
||||
Box::new(self.values.iter().copied())
|
||||
}
|
||||
|
||||
fn min_value(&self) -> T {
|
||||
self.min_value
|
||||
}
|
||||
|
||||
fn max_value(&self) -> T {
|
||||
self.max_value
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.values.len() as u32
|
||||
}
|
||||
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
output.copy_from_slice(&self.values[start as usize..][..output.len()])
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, T: Copy + PartialOrd + Default, V> From<&'a V> for VecColumn<'a, T>
|
||||
where V: AsRef<[T]> + ?Sized
|
||||
{
|
||||
fn from(values: &'a V) -> Self {
|
||||
let values = values.as_ref();
|
||||
let (min_value, max_value) = minmax(values.iter().copied()).unwrap_or_default();
|
||||
Self {
|
||||
values,
|
||||
min_value,
|
||||
max_value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct MonotonicMappingColumn<C, T, Input> {
|
||||
from_column: C,
|
||||
monotonic_mapping: T,
|
||||
_phantom: PhantomData<Input>,
|
||||
}
|
||||
|
||||
/// Creates a view of a column transformed by a strictly monotonic mapping. See
|
||||
/// [`StrictlyMonotonicFn`].
|
||||
///
|
||||
/// E.g. apply a gcd monotonic_mapping([100, 200, 300]) == [1, 2, 3]
|
||||
/// monotonic_mapping.mapping() is expected to be injective, and we should always have
|
||||
/// monotonic_mapping.inverse(monotonic_mapping.mapping(el)) == el
|
||||
///
|
||||
/// The inverse of the mapping is required for:
|
||||
/// `fn get_positions_for_value_range(&self, range: RangeInclusive<T>) -> Vec<u64> `
|
||||
/// The user provides the original value range and we need to monotonic map them in the same way the
|
||||
/// serialization does before calling the underlying column.
|
||||
///
|
||||
/// Note that when opening a codec, the monotonic_mapping should be the inverse of the mapping
|
||||
/// during serialization. And therefore the monotonic_mapping_inv when opening is the same as
|
||||
/// monotonic_mapping during serialization.
|
||||
pub fn monotonic_map_column<C, T, Input, Output>(
|
||||
from_column: C,
|
||||
monotonic_mapping: T,
|
||||
) -> impl ColumnValues<Output>
|
||||
where
|
||||
C: ColumnValues<Input>,
|
||||
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
|
||||
Input: PartialOrd + Debug + Send + Sync + Clone,
|
||||
Output: PartialOrd + Debug + Send + Sync + Clone,
|
||||
{
|
||||
MonotonicMappingColumn {
|
||||
from_column,
|
||||
monotonic_mapping,
|
||||
_phantom: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
impl<C, T, Input, Output> ColumnValues<Output> for MonotonicMappingColumn<C, T, Input>
|
||||
where
|
||||
C: ColumnValues<Input>,
|
||||
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
|
||||
Input: PartialOrd + Send + Debug + Sync + Clone,
|
||||
Output: PartialOrd + Send + Debug + Sync + Clone,
|
||||
{
|
||||
#[inline]
|
||||
fn get_val(&self, idx: u32) -> Output {
|
||||
let from_val = self.from_column.get_val(idx);
|
||||
self.monotonic_mapping.mapping(from_val)
|
||||
}
|
||||
|
||||
fn min_value(&self) -> Output {
|
||||
let from_min_value = self.from_column.min_value();
|
||||
self.monotonic_mapping.mapping(from_min_value)
|
||||
}
|
||||
|
||||
fn max_value(&self) -> Output {
|
||||
let from_max_value = self.from_column.max_value();
|
||||
self.monotonic_mapping.mapping(from_max_value)
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.from_column.num_vals()
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = Output> + '_> {
|
||||
Box::new(
|
||||
self.from_column
|
||||
.iter()
|
||||
.map(|el| self.monotonic_mapping.mapping(el)),
|
||||
)
|
||||
}
|
||||
|
||||
fn get_docids_for_value_range(
|
||||
&self,
|
||||
range: RangeInclusive<Output>,
|
||||
doc_id_range: Range<u32>,
|
||||
positions: &mut Vec<u32>,
|
||||
) {
|
||||
self.from_column.get_docids_for_value_range(
|
||||
self.monotonic_mapping.inverse(range.start().clone())
|
||||
..=self.monotonic_mapping.inverse(range.end().clone()),
|
||||
doc_id_range,
|
||||
positions,
|
||||
)
|
||||
}
|
||||
|
||||
// We voluntarily do not implement get_range as it yields a regression,
|
||||
// and we do not have any specialized implementation anyway.
|
||||
}
|
||||
|
||||
/// Wraps an iterator into a `Column`.
|
||||
pub struct IterColumn<T>(T);
|
||||
|
||||
impl<T> From<T> for IterColumn<T>
|
||||
where T: Iterator + Clone + ExactSizeIterator
|
||||
{
|
||||
fn from(iter: T) -> Self {
|
||||
IterColumn(iter)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> ColumnValues<T::Item> for IterColumn<T>
|
||||
where
|
||||
T: Iterator + Clone + ExactSizeIterator + Send + Sync,
|
||||
T::Item: PartialOrd + Debug,
|
||||
{
|
||||
fn get_val(&self, idx: u32) -> T::Item {
|
||||
self.0.clone().nth(idx as usize).unwrap()
|
||||
}
|
||||
|
||||
fn min_value(&self) -> T::Item {
|
||||
self.0.clone().next().unwrap()
|
||||
}
|
||||
|
||||
fn max_value(&self) -> T::Item {
|
||||
self.0.clone().last().unwrap()
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.0.len() as u32
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = T::Item> + '_> {
|
||||
Box::new(self.0.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::column_values::monotonic_mapping::{
|
||||
StrictlyMonotonicMappingInverter, StrictlyMonotonicMappingToInternalBaseval,
|
||||
StrictlyMonotonicMappingToInternalGCDBaseval,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_monotonic_mapping() {
|
||||
let vals = &[3u64, 5u64][..];
|
||||
let col = VecColumn::from(vals);
|
||||
let mapped = monotonic_map_column(col, StrictlyMonotonicMappingToInternalBaseval::new(2));
|
||||
assert_eq!(mapped.min_value(), 1u64);
|
||||
assert_eq!(mapped.max_value(), 3u64);
|
||||
assert_eq!(mapped.num_vals(), 2);
|
||||
assert_eq!(mapped.num_vals(), 2);
|
||||
assert_eq!(mapped.get_val(0), 1);
|
||||
assert_eq!(mapped.get_val(1), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_range_as_col() {
|
||||
let col = IterColumn::from(10..100);
|
||||
assert_eq!(col.num_vals(), 90);
|
||||
assert_eq!(col.max_value(), 99);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_monotonic_mapping_iter() {
|
||||
let vals: Vec<u64> = (10..110u64).map(|el| el * 10).collect();
|
||||
let col = VecColumn::from(&vals);
|
||||
let mapped = monotonic_map_column(
|
||||
col,
|
||||
StrictlyMonotonicMappingInverter::from(
|
||||
StrictlyMonotonicMappingToInternalGCDBaseval::new(10, 100),
|
||||
),
|
||||
);
|
||||
let val_i64s: Vec<u64> = mapped.iter().collect();
|
||||
for i in 0..100 {
|
||||
assert_eq!(val_i64s[i as usize], mapped.get_val(i));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_monotonic_mapping_get_range() {
|
||||
let vals: Vec<u64> = (0..100u64).map(|el| el * 10).collect();
|
||||
let col = VecColumn::from(&vals);
|
||||
let mapped = monotonic_map_column(
|
||||
col,
|
||||
StrictlyMonotonicMappingInverter::from(
|
||||
StrictlyMonotonicMappingToInternalGCDBaseval::new(10, 0),
|
||||
),
|
||||
);
|
||||
|
||||
assert_eq!(mapped.min_value(), 0u64);
|
||||
assert_eq!(mapped.max_value(), 9900u64);
|
||||
assert_eq!(mapped.num_vals(), 100);
|
||||
let val_u64s: Vec<u64> = mapped.iter().collect();
|
||||
assert_eq!(val_u64s.len(), 100);
|
||||
for i in 0..100 {
|
||||
assert_eq!(val_u64s[i as usize], mapped.get_val(i));
|
||||
assert_eq!(val_u64s[i as usize], vals[i as usize] * 10);
|
||||
}
|
||||
let mut buf = [0u64; 20];
|
||||
mapped.get_range(7, &mut buf[..]);
|
||||
assert_eq!(&val_u64s[7..][..20], &buf);
|
||||
}
|
||||
}
|
||||
@@ -17,15 +17,14 @@ use std::{
|
||||
ops::{Range, RangeInclusive},
|
||||
};
|
||||
|
||||
mod blank_range;
|
||||
mod build_compact_space;
|
||||
|
||||
use build_compact_space::get_compact_space;
|
||||
use common::{BinarySerializable, CountingWriter, OwnedBytes, VInt, VIntU128};
|
||||
use tantivy_bitpacker::{self, BitPacker, BitUnpacker};
|
||||
|
||||
use crate::column_values::compact_space::build_compact_space::get_compact_space;
|
||||
use crate::column_values::ColumnValues;
|
||||
use crate::RowId;
|
||||
|
||||
mod blank_range;
|
||||
mod build_compact_space;
|
||||
|
||||
/// The cost per blank is quite hard actually, since blanks are delta encoded, the actual cost of
|
||||
/// blanks depends on the number of blanks.
|
||||
@@ -159,30 +158,23 @@ impl CompactSpace {
|
||||
pub struct CompactSpaceCompressor {
|
||||
params: IPCodecParams,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct IPCodecParams {
|
||||
compact_space: CompactSpace,
|
||||
bit_unpacker: BitUnpacker,
|
||||
min_value: u128,
|
||||
max_value: u128,
|
||||
num_vals: RowId,
|
||||
num_vals: u32,
|
||||
num_bits: u8,
|
||||
}
|
||||
|
||||
impl CompactSpaceCompressor {
|
||||
pub fn num_vals(&self) -> RowId {
|
||||
self.params.num_vals
|
||||
}
|
||||
|
||||
/// Taking the vals as Vec may cost a lot of memory. It is used to sort the vals.
|
||||
pub fn train_from(iter: impl Iterator<Item = u128>) -> Self {
|
||||
pub fn train_from(iter: impl Iterator<Item = u128>, num_vals: u32) -> Self {
|
||||
let mut values_sorted = BTreeSet::new();
|
||||
let mut total_num_values = 0u32;
|
||||
for val in iter {
|
||||
total_num_values += 1u32;
|
||||
values_sorted.insert(val);
|
||||
}
|
||||
values_sorted.extend(iter);
|
||||
let total_num_values = num_vals;
|
||||
|
||||
let compact_space =
|
||||
get_compact_space(&values_sorted, total_num_values, COST_PER_BLANK_IN_BITS);
|
||||
let amplitude_compact_space = compact_space.amplitude_compact_space();
|
||||
@@ -313,7 +305,7 @@ impl ColumnValues<u128> for CompactSpaceDecompressor {
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn get_row_ids_for_value_range(
|
||||
fn get_docids_for_value_range(
|
||||
&self,
|
||||
value_range: RangeInclusive<u128>,
|
||||
positions_range: Range<u32>,
|
||||
@@ -458,352 +450,364 @@ impl CompactSpaceDecompressor {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use itertools::Itertools;
|
||||
|
||||
use super::*;
|
||||
use crate::column_values::u128_based::U128Header;
|
||||
use crate::column_values::{open_u128_mapped, serialize_column_values_u128};
|
||||
|
||||
#[test]
|
||||
fn compact_space_test() {
|
||||
let ips = &[
|
||||
2u128, 4u128, 1000, 1001, 1002, 1003, 1004, 1005, 1008, 1010, 1012, 1260,
|
||||
]
|
||||
.into_iter()
|
||||
.collect();
|
||||
let compact_space = get_compact_space(ips, ips.len() as u32, 11);
|
||||
let amplitude = compact_space.amplitude_compact_space();
|
||||
assert_eq!(amplitude, 17);
|
||||
assert_eq!(1, compact_space.u128_to_compact(2).unwrap());
|
||||
assert_eq!(2, compact_space.u128_to_compact(3).unwrap());
|
||||
assert_eq!(compact_space.u128_to_compact(100).unwrap_err(), 1);
|
||||
|
||||
for (num1, num2) in (0..3).tuple_windows() {
|
||||
assert_eq!(
|
||||
compact_space.get_range_mapping(num1).compact_end() + 1,
|
||||
compact_space.get_range_mapping(num2).compact_start
|
||||
);
|
||||
}
|
||||
|
||||
let mut output: Vec<u8> = Vec::new();
|
||||
compact_space.serialize(&mut output).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
compact_space,
|
||||
CompactSpace::deserialize(&mut &output[..]).unwrap()
|
||||
);
|
||||
|
||||
for ip in ips {
|
||||
let compact = compact_space.u128_to_compact(*ip).unwrap();
|
||||
assert_eq!(compact_space.compact_to_u128(compact), *ip);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compact_space_amplitude_test() {
|
||||
let ips = &[100000u128, 1000000].into_iter().collect();
|
||||
let compact_space = get_compact_space(ips, ips.len() as u32, 1);
|
||||
let amplitude = compact_space.amplitude_compact_space();
|
||||
assert_eq!(amplitude, 2);
|
||||
}
|
||||
|
||||
fn test_all(mut data: OwnedBytes, expected: &[u128]) {
|
||||
let _header = U128Header::deserialize(&mut data);
|
||||
let decompressor = CompactSpaceDecompressor::open(data).unwrap();
|
||||
for (idx, expected_val) in expected.iter().cloned().enumerate() {
|
||||
let val = decompressor.get(idx as u32);
|
||||
assert_eq!(val, expected_val);
|
||||
|
||||
let test_range = |range: RangeInclusive<u128>| {
|
||||
let expected_positions = expected
|
||||
.iter()
|
||||
.positions(|val| range.contains(val))
|
||||
.map(|pos| pos as u32)
|
||||
.collect::<Vec<_>>();
|
||||
let mut positions = Vec::new();
|
||||
decompressor.get_positions_for_value_range(
|
||||
range,
|
||||
0..decompressor.num_vals(),
|
||||
&mut positions,
|
||||
);
|
||||
assert_eq!(positions, expected_positions);
|
||||
};
|
||||
|
||||
test_range(expected_val.saturating_sub(1)..=expected_val);
|
||||
test_range(expected_val..=expected_val);
|
||||
test_range(expected_val..=expected_val.saturating_add(1));
|
||||
test_range(expected_val.saturating_sub(1)..=expected_val.saturating_add(1));
|
||||
}
|
||||
}
|
||||
|
||||
fn test_aux_vals(u128_vals: &[u128]) -> OwnedBytes {
|
||||
let mut out = Vec::new();
|
||||
serialize_column_values_u128(&u128_vals, &mut out).unwrap();
|
||||
let data = OwnedBytes::new(out);
|
||||
test_all(data.clone(), u128_vals);
|
||||
data
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_range_1() {
|
||||
let vals = &[
|
||||
1u128,
|
||||
100u128,
|
||||
3u128,
|
||||
99999u128,
|
||||
100000u128,
|
||||
100001u128,
|
||||
4_000_211_221u128,
|
||||
4_000_211_222u128,
|
||||
333u128,
|
||||
];
|
||||
let mut data = test_aux_vals(vals);
|
||||
|
||||
let _header = U128Header::deserialize(&mut data);
|
||||
let decomp = CompactSpaceDecompressor::open(data).unwrap();
|
||||
let complete_range = 0..vals.len() as u32;
|
||||
for (pos, val) in vals.iter().enumerate() {
|
||||
let val = *val;
|
||||
let pos = pos as u32;
|
||||
let mut positions = Vec::new();
|
||||
decomp.get_positions_for_value_range(val..=val, pos..pos + 1, &mut positions);
|
||||
assert_eq!(positions, vec![pos]);
|
||||
}
|
||||
|
||||
// handle docid range out of bounds
|
||||
let positions: Vec<u32> = get_positions_for_value_range_helper(&decomp, 0..=1, 1..u32::MAX);
|
||||
assert!(positions.is_empty());
|
||||
|
||||
let positions =
|
||||
get_positions_for_value_range_helper(&decomp, 0..=1, complete_range.clone());
|
||||
assert_eq!(positions, vec![0]);
|
||||
let positions =
|
||||
get_positions_for_value_range_helper(&decomp, 0..=2, complete_range.clone());
|
||||
assert_eq!(positions, vec![0]);
|
||||
let positions =
|
||||
get_positions_for_value_range_helper(&decomp, 0..=3, complete_range.clone());
|
||||
assert_eq!(positions, vec![0, 2]);
|
||||
assert_eq!(
|
||||
get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
99999u128..=99999u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
vec![3]
|
||||
);
|
||||
assert_eq!(
|
||||
get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
99999u128..=100000u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
vec![3, 4]
|
||||
);
|
||||
assert_eq!(
|
||||
get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
99998u128..=100000u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
vec![3, 4]
|
||||
);
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
99998u128..=99999u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
&[3]
|
||||
);
|
||||
assert!(get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
99998u128..=99998u128,
|
||||
complete_range.clone()
|
||||
)
|
||||
.is_empty());
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
333u128..=333u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
&[8]
|
||||
);
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
332u128..=333u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
&[8]
|
||||
);
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
332u128..=334u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
&[8]
|
||||
);
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
333u128..=334u128,
|
||||
complete_range.clone()
|
||||
),
|
||||
&[8]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(
|
||||
&decomp,
|
||||
4_000_211_221u128..=5_000_000_000u128,
|
||||
complete_range
|
||||
),
|
||||
&[6, 7]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty() {
|
||||
let vals = &[];
|
||||
let data = test_aux_vals(vals);
|
||||
let _decomp = CompactSpaceDecompressor::open(data).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_range_2() {
|
||||
let vals = &[
|
||||
100u128,
|
||||
99999u128,
|
||||
100000u128,
|
||||
100001u128,
|
||||
4_000_211_221u128,
|
||||
4_000_211_222u128,
|
||||
333u128,
|
||||
];
|
||||
let mut data = test_aux_vals(vals);
|
||||
let _header = U128Header::deserialize(&mut data);
|
||||
let decomp = CompactSpaceDecompressor::open(data).unwrap();
|
||||
let complete_range = 0..vals.len() as u32;
|
||||
assert!(
|
||||
&get_positions_for_value_range_helper(&decomp, 0..=5, complete_range.clone())
|
||||
.is_empty(),
|
||||
);
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(&decomp, 0..=100, complete_range.clone()),
|
||||
&[0]
|
||||
);
|
||||
assert_eq!(
|
||||
&get_positions_for_value_range_helper(&decomp, 0..=105, complete_range),
|
||||
&[0]
|
||||
);
|
||||
}
|
||||
|
||||
fn get_positions_for_value_range_helper<C: ColumnValues<T> + ?Sized, T: PartialOrd>(
|
||||
column: &C,
|
||||
value_range: RangeInclusive<T>,
|
||||
doc_id_range: Range<u32>,
|
||||
) -> Vec<u32> {
|
||||
let mut positions = Vec::new();
|
||||
column.get_row_ids_for_value_range(value_range, doc_id_range, &mut positions);
|
||||
positions
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_range_3() {
|
||||
let vals = &[
|
||||
200u128,
|
||||
201,
|
||||
202,
|
||||
203,
|
||||
204,
|
||||
204,
|
||||
206,
|
||||
207,
|
||||
208,
|
||||
209,
|
||||
210,
|
||||
1_000_000,
|
||||
5_000_000_000,
|
||||
];
|
||||
let mut out = Vec::new();
|
||||
serialize_column_values_u128(&&vals[..], &mut out).unwrap();
|
||||
let decomp = open_u128_mapped(OwnedBytes::new(out)).unwrap();
|
||||
let complete_range = 0..vals.len() as u32;
|
||||
|
||||
assert_eq!(
|
||||
get_positions_for_value_range_helper(&*decomp, 199..=200, complete_range.clone()),
|
||||
vec![0]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
get_positions_for_value_range_helper(&*decomp, 199..=201, complete_range.clone()),
|
||||
vec![0, 1]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
get_positions_for_value_range_helper(&*decomp, 200..=200, complete_range.clone()),
|
||||
vec![0]
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
get_positions_for_value_range_helper(&*decomp, 1_000_000..=1_000_000, complete_range),
|
||||
vec![11]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug1() {
|
||||
let vals = &[9223372036854775806];
|
||||
let _data = test_aux_vals(vals);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug2() {
|
||||
let vals = &[340282366920938463463374607431768211455u128];
|
||||
let _data = test_aux_vals(vals);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug3() {
|
||||
let vals = &[340282366920938463463374607431768211454];
|
||||
let _data = test_aux_vals(vals);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bug4() {
|
||||
let vals = &[340282366920938463463374607431768211455, 0];
|
||||
let _data = test_aux_vals(vals);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_first_large_gaps() {
|
||||
let vals = &[1_000_000_000u128; 100];
|
||||
let _data = test_aux_vals(vals);
|
||||
}
|
||||
|
||||
use proptest::prelude::*;
|
||||
|
||||
fn num_strategy() -> impl Strategy<Value = u128> {
|
||||
prop_oneof![
|
||||
1 => prop::num::u128::ANY.prop_map(|num| u128::MAX - (num % 10) ),
|
||||
1 => prop::num::u128::ANY.prop_map(|num| i64::MAX as u128 + 5 - (num % 10) ),
|
||||
1 => prop::num::u128::ANY.prop_map(|num| i128::MAX as u128 + 5 - (num % 10) ),
|
||||
1 => prop::num::u128::ANY.prop_map(|num| num % 10 ),
|
||||
20 => prop::num::u128::ANY,
|
||||
]
|
||||
}
|
||||
|
||||
proptest! {
|
||||
#![proptest_config(ProptestConfig::with_cases(10))]
|
||||
|
||||
#[test]
|
||||
fn compress_decompress_random(vals in proptest::collection::vec(num_strategy() , 1..1000)) {
|
||||
let _data = test_aux_vals(&vals);
|
||||
}
|
||||
}
|
||||
}
|
||||
// TODO reenable what can be reenabled.
|
||||
// #[cfg(test)]
|
||||
// mod tests {
|
||||
//
|
||||
// use super::*;
|
||||
// use crate::column::format_version::read_format_version;
|
||||
// use crate::column::column_footer::read_null_index_footer;
|
||||
// use crate::column::serialize::U128Header;
|
||||
// use crate::column::{open_u128, serialize_u128};
|
||||
//
|
||||
// #[test]
|
||||
// fn compact_space_test() {
|
||||
// let ips = &[
|
||||
// 2u128, 4u128, 1000, 1001, 1002, 1003, 1004, 1005, 1008, 1010, 1012, 1260,
|
||||
// ]
|
||||
// .into_iter()
|
||||
// .collect();
|
||||
// let compact_space = get_compact_space(ips, ips.len() as u32, 11);
|
||||
// let amplitude = compact_space.amplitude_compact_space();
|
||||
// assert_eq!(amplitude, 17);
|
||||
// assert_eq!(1, compact_space.u128_to_compact(2).unwrap());
|
||||
// assert_eq!(2, compact_space.u128_to_compact(3).unwrap());
|
||||
// assert_eq!(compact_space.u128_to_compact(100).unwrap_err(), 1);
|
||||
//
|
||||
// for (num1, num2) in (0..3).tuple_windows() {
|
||||
// assert_eq!(
|
||||
// compact_space.get_range_mapping(num1).compact_end() + 1,
|
||||
// compact_space.get_range_mapping(num2).compact_start
|
||||
// );
|
||||
// }
|
||||
//
|
||||
// let mut output: Vec<u8> = Vec::new();
|
||||
// compact_space.serialize(&mut output).unwrap();
|
||||
//
|
||||
// assert_eq!(
|
||||
// compact_space,
|
||||
// CompactSpace::deserialize(&mut &output[..]).unwrap()
|
||||
// );
|
||||
//
|
||||
// for ip in ips {
|
||||
// let compact = compact_space.u128_to_compact(*ip).unwrap();
|
||||
// assert_eq!(compact_space.compact_to_u128(compact), *ip);
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn compact_space_amplitude_test() {
|
||||
// let ips = &[100000u128, 1000000].into_iter().collect();
|
||||
// let compact_space = get_compact_space(ips, ips.len() as u32, 1);
|
||||
// let amplitude = compact_space.amplitude_compact_space();
|
||||
// assert_eq!(amplitude, 2);
|
||||
// }
|
||||
//
|
||||
// fn test_all(mut data: OwnedBytes, expected: &[u128]) {
|
||||
// let _header = U128Header::deserialize(&mut data);
|
||||
// let decompressor = CompactSpaceDecompressor::open(data).unwrap();
|
||||
// for (idx, expected_val) in expected.iter().cloned().enumerate() {
|
||||
// let val = decompressor.get(idx as u32);
|
||||
// assert_eq!(val, expected_val);
|
||||
//
|
||||
// let test_range = |range: RangeInclusive<u128>| {
|
||||
// let expected_positions = expected
|
||||
// .iter()
|
||||
// .positions(|val| range.contains(val))
|
||||
// .map(|pos| pos as u32)
|
||||
// .collect::<Vec<_>>();
|
||||
// let mut positions = Vec::new();
|
||||
// decompressor.get_positions_for_value_range(
|
||||
// range,
|
||||
// 0..decompressor.num_vals(),
|
||||
// &mut positions,
|
||||
// );
|
||||
// assert_eq!(positions, expected_positions);
|
||||
// };
|
||||
//
|
||||
// test_range(expected_val.saturating_sub(1)..=expected_val);
|
||||
// test_range(expected_val..=expected_val);
|
||||
// test_range(expected_val..=expected_val.saturating_add(1));
|
||||
// test_range(expected_val.saturating_sub(1)..=expected_val.saturating_add(1));
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// fn test_aux_vals(u128_vals: &[u128]) -> OwnedBytes {
|
||||
// let mut out = Vec::new();
|
||||
// serialize_u128(
|
||||
// || u128_vals.iter().cloned(),
|
||||
// u128_vals.len() as u32,
|
||||
// &mut out,
|
||||
// )
|
||||
// .unwrap();
|
||||
//
|
||||
// let data = OwnedBytes::new(out);
|
||||
// let (data, _format_version) = read_format_version(data).unwrap();
|
||||
// let (data, _null_index_footer) = read_null_index_footer(data).unwrap();
|
||||
// test_all(data.clone(), u128_vals);
|
||||
//
|
||||
// data
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_range_1() {
|
||||
// let vals = &[
|
||||
// 1u128,
|
||||
// 100u128,
|
||||
// 3u128,
|
||||
// 99999u128,
|
||||
// 100000u128,
|
||||
// 100001u128,
|
||||
// 4_000_211_221u128,
|
||||
// 4_000_211_222u128,
|
||||
// 333u128,
|
||||
// ];
|
||||
// let mut data = test_aux_vals(vals);
|
||||
//
|
||||
// let _header = U128Header::deserialize(&mut data);
|
||||
// let decomp = CompactSpaceDecompressor::open(data).unwrap();
|
||||
// let complete_range = 0..vals.len() as u32;
|
||||
// for (pos, val) in vals.iter().enumerate() {
|
||||
// let val = *val;
|
||||
// let pos = pos as u32;
|
||||
// let mut positions = Vec::new();
|
||||
// decomp.get_positions_for_value_range(val..=val, pos..pos + 1, &mut positions);
|
||||
// assert_eq!(positions, vec![pos]);
|
||||
// }
|
||||
//
|
||||
// handle docid range out of bounds
|
||||
// let positions: Vec<u32> = get_positions_for_value_range_helper(&decomp, 0..=1, 1..u32::MAX);
|
||||
// assert!(positions.is_empty());
|
||||
//
|
||||
// let positions =
|
||||
// get_positions_for_value_range_helper(&decomp, 0..=1, complete_range.clone());
|
||||
// assert_eq!(positions, vec![0]);
|
||||
// let positions =
|
||||
// get_positions_for_value_range_helper(&decomp, 0..=2, complete_range.clone());
|
||||
// assert_eq!(positions, vec![0]);
|
||||
// let positions =
|
||||
// get_positions_for_value_range_helper(&decomp, 0..=3, complete_range.clone());
|
||||
// assert_eq!(positions, vec![0, 2]);
|
||||
// assert_eq!(
|
||||
// get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 99999u128..=99999u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// vec![3]
|
||||
// );
|
||||
// assert_eq!(
|
||||
// get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 99999u128..=100000u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// vec![3, 4]
|
||||
// );
|
||||
// assert_eq!(
|
||||
// get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 99998u128..=100000u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// vec![3, 4]
|
||||
// );
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 99998u128..=99999u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// &[3]
|
||||
// );
|
||||
// assert!(get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 99998u128..=99998u128,
|
||||
// complete_range.clone()
|
||||
// )
|
||||
// .is_empty());
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 333u128..=333u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// &[8]
|
||||
// );
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 332u128..=333u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// &[8]
|
||||
// );
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 332u128..=334u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// &[8]
|
||||
// );
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 333u128..=334u128,
|
||||
// complete_range.clone()
|
||||
// ),
|
||||
// &[8]
|
||||
// );
|
||||
//
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(
|
||||
// &decomp,
|
||||
// 4_000_211_221u128..=5_000_000_000u128,
|
||||
// complete_range
|
||||
// ),
|
||||
// &[6, 7]
|
||||
// );
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_empty() {
|
||||
// let vals = &[];
|
||||
// let data = test_aux_vals(vals);
|
||||
// let _decomp = CompactSpaceDecompressor::open(data).unwrap();
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_range_2() {
|
||||
// let vals = &[
|
||||
// 100u128,
|
||||
// 99999u128,
|
||||
// 100000u128,
|
||||
// 100001u128,
|
||||
// 4_000_211_221u128,
|
||||
// 4_000_211_222u128,
|
||||
// 333u128,
|
||||
// ];
|
||||
// let mut data = test_aux_vals(vals);
|
||||
// let _header = U128Header::deserialize(&mut data);
|
||||
// let decomp = CompactSpaceDecompressor::open(data).unwrap();
|
||||
// let complete_range = 0..vals.len() as u32;
|
||||
// assert!(
|
||||
// &get_positions_for_value_range_helper(&decomp, 0..=5, complete_range.clone())
|
||||
// .is_empty(),
|
||||
// );
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(&decomp, 0..=100, complete_range.clone()),
|
||||
// &[0]
|
||||
// );
|
||||
// assert_eq!(
|
||||
// &get_positions_for_value_range_helper(&decomp, 0..=105, complete_range),
|
||||
// &[0]
|
||||
// );
|
||||
// }
|
||||
//
|
||||
// fn get_positions_for_value_range_helper<C: Column<T> + ?Sized, T: PartialOrd>(
|
||||
// column: &C,
|
||||
// value_range: RangeInclusive<T>,
|
||||
// doc_id_range: Range<u32>,
|
||||
// ) -> Vec<u32> {
|
||||
// let mut positions = Vec::new();
|
||||
// column.get_docids_for_value_range(value_range, doc_id_range, &mut positions);
|
||||
// positions
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_range_3() {
|
||||
// let vals = &[
|
||||
// 200u128,
|
||||
// 201,
|
||||
// 202,
|
||||
// 203,
|
||||
// 204,
|
||||
// 204,
|
||||
// 206,
|
||||
// 207,
|
||||
// 208,
|
||||
// 209,
|
||||
// 210,
|
||||
// 1_000_000,
|
||||
// 5_000_000_000,
|
||||
// ];
|
||||
// let mut out = Vec::new();
|
||||
// serialize_u128(|| vals.iter().cloned(), vals.len() as u32, &mut out).unwrap();
|
||||
// let decomp = open_u128::<u128>(OwnedBytes::new(out)).unwrap();
|
||||
// let complete_range = 0..vals.len() as u32;
|
||||
//
|
||||
// assert_eq!(
|
||||
// get_positions_for_value_range_helper(&*decomp, 199..=200, complete_range.clone()),
|
||||
// vec![0]
|
||||
// );
|
||||
//
|
||||
// assert_eq!(
|
||||
// get_positions_for_value_range_helper(&*decomp, 199..=201, complete_range.clone()),
|
||||
// vec![0, 1]
|
||||
// );
|
||||
//
|
||||
// assert_eq!(
|
||||
// get_positions_for_value_range_helper(&*decomp, 200..=200, complete_range.clone()),
|
||||
// vec![0]
|
||||
// );
|
||||
//
|
||||
// assert_eq!(
|
||||
// get_positions_for_value_range_helper(&*decomp, 1_000_000..=1_000_000, complete_range),
|
||||
// vec![11]
|
||||
// );
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_bug1() {
|
||||
// let vals = &[9223372036854775806];
|
||||
// let _data = test_aux_vals(vals);
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_bug2() {
|
||||
// let vals = &[340282366920938463463374607431768211455u128];
|
||||
// let _data = test_aux_vals(vals);
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_bug3() {
|
||||
// let vals = &[340282366920938463463374607431768211454];
|
||||
// let _data = test_aux_vals(vals);
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_bug4() {
|
||||
// let vals = &[340282366920938463463374607431768211455, 0];
|
||||
// let _data = test_aux_vals(vals);
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn test_first_large_gaps() {
|
||||
// let vals = &[1_000_000_000u128; 100];
|
||||
// let _data = test_aux_vals(vals);
|
||||
// }
|
||||
// use itertools::Itertools;
|
||||
// use proptest::prelude::*;
|
||||
//
|
||||
// fn num_strategy() -> impl Strategy<Value = u128> {
|
||||
// prop_oneof![
|
||||
// 1 => prop::num::u128::ANY.prop_map(|num| u128::MAX - (num % 10) ),
|
||||
// 1 => prop::num::u128::ANY.prop_map(|num| i64::MAX as u128 + 5 - (num % 10) ),
|
||||
// 1 => prop::num::u128::ANY.prop_map(|num| i128::MAX as u128 + 5 - (num % 10) ),
|
||||
// 1 => prop::num::u128::ANY.prop_map(|num| num % 10 ),
|
||||
// 20 => prop::num::u128::ANY,
|
||||
// ]
|
||||
// }
|
||||
//
|
||||
// proptest! {
|
||||
// #![proptest_config(ProptestConfig::with_cases(10))]
|
||||
//
|
||||
// #[test]
|
||||
// fn compress_decompress_random(vals in proptest::collection::vec(num_strategy()
|
||||
// , 1..1000)) {
|
||||
// let _data = test_aux_vals(&vals);
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
222
columnar/src/column_values/main.rs
Normal file
222
columnar/src/column_values/main.rs
Normal file
@@ -0,0 +1,222 @@
|
||||
#[macro_use]
|
||||
extern crate prettytable;
|
||||
use std::collections::HashSet;
|
||||
use std::env;
|
||||
use std::io::BufRead;
|
||||
use std::net::{IpAddr, Ipv6Addr};
|
||||
use std::str::FromStr;
|
||||
|
||||
use common::OwnedBytes;
|
||||
use fastfield_codecs::{open_u128, serialize_u128, Column, FastFieldCodecType, VecColumn};
|
||||
use itertools::Itertools;
|
||||
use measure_time::print_time;
|
||||
use prettytable::{Cell, Row, Table};
|
||||
|
||||
fn print_set_stats(ip_addrs: &[u128]) {
|
||||
println!("NumIps\t{}", ip_addrs.len());
|
||||
let ip_addr_set: HashSet<u128> = ip_addrs.iter().cloned().collect();
|
||||
println!("NumUniqueIps\t{}", ip_addr_set.len());
|
||||
let ratio_unique = ip_addr_set.len() as f64 / ip_addrs.len() as f64;
|
||||
println!("RatioUniqueOverTotal\t{ratio_unique:.4}");
|
||||
|
||||
// histogram
|
||||
let mut ip_addrs = ip_addrs.to_vec();
|
||||
ip_addrs.sort();
|
||||
let mut cnts: Vec<usize> = ip_addrs
|
||||
.into_iter()
|
||||
.dedup_with_count()
|
||||
.map(|(cnt, _)| cnt)
|
||||
.collect();
|
||||
cnts.sort();
|
||||
|
||||
let top_256_cnt: usize = cnts.iter().rev().take(256).sum();
|
||||
let top_128_cnt: usize = cnts.iter().rev().take(128).sum();
|
||||
let top_64_cnt: usize = cnts.iter().rev().take(64).sum();
|
||||
let top_8_cnt: usize = cnts.iter().rev().take(8).sum();
|
||||
let total: usize = cnts.iter().sum();
|
||||
|
||||
println!("{}", total);
|
||||
println!("{}", top_256_cnt);
|
||||
println!("{}", top_128_cnt);
|
||||
println!("Percentage Top8 {:02}", top_8_cnt as f32 / total as f32);
|
||||
println!("Percentage Top64 {:02}", top_64_cnt as f32 / total as f32);
|
||||
println!("Percentage Top128 {:02}", top_128_cnt as f32 / total as f32);
|
||||
println!("Percentage Top256 {:02}", top_256_cnt as f32 / total as f32);
|
||||
|
||||
let mut cnts: Vec<(usize, usize)> = cnts.into_iter().dedup_with_count().collect();
|
||||
cnts.sort_by(|a, b| {
|
||||
if a.1 == b.1 {
|
||||
a.0.cmp(&b.0)
|
||||
} else {
|
||||
b.1.cmp(&a.1)
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn ip_dataset() -> Vec<u128> {
|
||||
let mut ip_addr_v4 = 0;
|
||||
|
||||
let stdin = std::io::stdin();
|
||||
let ip_addrs: Vec<u128> = stdin
|
||||
.lock()
|
||||
.lines()
|
||||
.flat_map(|line| {
|
||||
let line = line.unwrap();
|
||||
let line = line.trim();
|
||||
let ip_addr = IpAddr::from_str(line.trim()).ok()?;
|
||||
if ip_addr.is_ipv4() {
|
||||
ip_addr_v4 += 1;
|
||||
}
|
||||
let ip_addr_v6: Ipv6Addr = match ip_addr {
|
||||
IpAddr::V4(v4) => v4.to_ipv6_mapped(),
|
||||
IpAddr::V6(v6) => v6,
|
||||
};
|
||||
Some(ip_addr_v6)
|
||||
})
|
||||
.map(|ip_v6| u128::from_be_bytes(ip_v6.octets()))
|
||||
.collect();
|
||||
|
||||
println!("IpAddrsAny\t{}", ip_addrs.len());
|
||||
println!("IpAddrsV4\t{}", ip_addr_v4);
|
||||
|
||||
ip_addrs
|
||||
}
|
||||
|
||||
fn bench_ip() {
|
||||
let dataset = ip_dataset();
|
||||
print_set_stats(&dataset);
|
||||
|
||||
// Chunks
|
||||
{
|
||||
let mut data = vec![];
|
||||
for dataset in dataset.chunks(500_000) {
|
||||
serialize_u128(|| dataset.iter().cloned(), dataset.len() as u32, &mut data).unwrap();
|
||||
}
|
||||
let compression = data.len() as f64 / (dataset.len() * 16) as f64;
|
||||
println!("Compression 50_000 chunks {:.4}", compression);
|
||||
println!(
|
||||
"Num Bits per elem {:.2}",
|
||||
(data.len() * 8) as f32 / dataset.len() as f32
|
||||
);
|
||||
}
|
||||
|
||||
let mut data = vec![];
|
||||
{
|
||||
print_time!("creation");
|
||||
serialize_u128(|| dataset.iter().cloned(), dataset.len() as u32, &mut data).unwrap();
|
||||
}
|
||||
|
||||
let compression = data.len() as f64 / (dataset.len() * 16) as f64;
|
||||
println!("Compression {:.2}", compression);
|
||||
println!(
|
||||
"Num Bits per elem {:.2}",
|
||||
(data.len() * 8) as f32 / dataset.len() as f32
|
||||
);
|
||||
|
||||
let decompressor = open_u128::<u128>(OwnedBytes::new(data)).unwrap();
|
||||
// Sample some ranges
|
||||
let mut doc_values = Vec::new();
|
||||
for value in dataset.iter().take(1110).skip(1100).cloned() {
|
||||
doc_values.clear();
|
||||
print_time!("get range");
|
||||
decompressor.get_docids_for_value_range(
|
||||
value..=value,
|
||||
0..decompressor.num_vals(),
|
||||
&mut doc_values,
|
||||
);
|
||||
println!("{:?}", doc_values.len());
|
||||
}
|
||||
}
|
||||
|
||||
fn main() {
|
||||
if env::args().nth(1).unwrap() == "bench_ip" {
|
||||
bench_ip();
|
||||
return;
|
||||
}
|
||||
|
||||
let mut table = Table::new();
|
||||
|
||||
// Add a row per time
|
||||
table.add_row(row!["", "Compression Ratio", "Compression Estimation"]);
|
||||
|
||||
for (data, data_set_name) in get_codec_test_data_sets() {
|
||||
let results: Vec<(f32, f32, FastFieldCodecType)> = [
|
||||
serialize_with_codec(&data, FastFieldCodecType::Bitpacked),
|
||||
serialize_with_codec(&data, FastFieldCodecType::Linear),
|
||||
serialize_with_codec(&data, FastFieldCodecType::BlockwiseLinear),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect();
|
||||
let best_compression_ratio_codec = results
|
||||
.iter()
|
||||
.min_by(|&res1, &res2| res1.partial_cmp(res2).unwrap())
|
||||
.cloned()
|
||||
.unwrap();
|
||||
|
||||
table.add_row(Row::new(vec![Cell::new(data_set_name).style_spec("Bbb")]));
|
||||
for (est, comp, codec_type) in results {
|
||||
let est_cell = est.to_string();
|
||||
let ratio_cell = comp.to_string();
|
||||
let style = if comp == best_compression_ratio_codec.1 {
|
||||
"Fb"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
table.add_row(Row::new(vec![
|
||||
Cell::new(&format!("{codec_type:?}")).style_spec("bFg"),
|
||||
Cell::new(&ratio_cell).style_spec(style),
|
||||
Cell::new(&est_cell).style_spec(""),
|
||||
]));
|
||||
}
|
||||
}
|
||||
|
||||
table.printstd();
|
||||
}
|
||||
|
||||
pub fn get_codec_test_data_sets() -> Vec<(Vec<u64>, &'static str)> {
|
||||
let mut data_and_names = vec![];
|
||||
|
||||
let data = (1000..=200_000_u64).collect::<Vec<_>>();
|
||||
data_and_names.push((data, "Autoincrement"));
|
||||
|
||||
let mut current_cumulative = 0;
|
||||
let data = (1..=200_000_u64)
|
||||
.map(|num| {
|
||||
let num = (num as f32 + num as f32).log10() as u64;
|
||||
current_cumulative += num;
|
||||
current_cumulative
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
// let data = (1..=200000_u64).map(|num| num + num).collect::<Vec<_>>();
|
||||
data_and_names.push((data, "Monotonically increasing concave"));
|
||||
|
||||
let mut current_cumulative = 0;
|
||||
let data = (1..=200_000_u64)
|
||||
.map(|num| {
|
||||
let num = (200_000.0 - num as f32).log10() as u64;
|
||||
current_cumulative += num;
|
||||
current_cumulative
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
data_and_names.push((data, "Monotonically increasing convex"));
|
||||
|
||||
let data = (1000..=200_000_u64)
|
||||
.map(|num| num + rand::random::<u8>() as u64)
|
||||
.collect::<Vec<_>>();
|
||||
data_and_names.push((data, "Almost monotonically increasing"));
|
||||
|
||||
data_and_names
|
||||
}
|
||||
|
||||
pub fn serialize_with_codec(
|
||||
data: &[u64],
|
||||
codec_type: FastFieldCodecType,
|
||||
) -> Option<(f32, f32, FastFieldCodecType)> {
|
||||
let col = VecColumn::from(data);
|
||||
let estimation = fastfield_codecs::estimate(&col, codec_type)?;
|
||||
let mut out = Vec::new();
|
||||
fastfield_codecs::serialize(&col, &mut out, &[codec_type]).ok()?;
|
||||
let actual_compression = out.len() as f32 / (col.num_vals() * 8) as f32;
|
||||
Some((estimation, actual_compression, codec_type))
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
use std::fmt::Debug;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{ColumnIndex, ColumnValues, MergeRowOrder};
|
||||
|
||||
pub(crate) struct MergedColumnValues<'a, T> {
|
||||
pub(crate) column_indexes: &'a [ColumnIndex],
|
||||
pub(crate) column_values: &'a [Option<Arc<dyn ColumnValues<T>>>],
|
||||
pub(crate) merge_row_order: &'a MergeRowOrder,
|
||||
}
|
||||
|
||||
impl<'a, T: Copy + PartialOrd + Debug> Iterable<T> for MergedColumnValues<'a, T> {
|
||||
fn boxed_iter(&self) -> Box<dyn Iterator<Item = T> + '_> {
|
||||
match self.merge_row_order {
|
||||
MergeRowOrder::Stack(_) => Box::new(
|
||||
self.column_values
|
||||
.iter()
|
||||
.flatten()
|
||||
.flat_map(|column_value| column_value.iter()),
|
||||
),
|
||||
MergeRowOrder::Shuffled(shuffle_merge_order) => Box::new(
|
||||
shuffle_merge_order
|
||||
.iter_new_to_old_row_addrs()
|
||||
.flat_map(|row_addr| {
|
||||
let column_index = &self.column_indexes[row_addr.segment_ord as usize];
|
||||
let column_values =
|
||||
self.column_values[row_addr.segment_ord as usize].as_ref()?;
|
||||
let value_range = column_index.value_row_ids(row_addr.row_id);
|
||||
Some((value_range, column_values))
|
||||
})
|
||||
.flat_map(|(value_range, column_values)| {
|
||||
value_range
|
||||
.into_iter()
|
||||
.map(|val| column_values.get_val(val))
|
||||
}),
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
#![warn(missing_docs)]
|
||||
#![cfg_attr(all(feature = "unstable", test), feature(test))]
|
||||
|
||||
//! # `fastfield_codecs`
|
||||
//!
|
||||
@@ -7,203 +8,206 @@
|
||||
//! - Monotonically map values to u64/u128
|
||||
|
||||
use std::fmt::Debug;
|
||||
use std::ops::{Range, RangeInclusive};
|
||||
use std::io;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::{BinarySerializable, OwnedBytes};
|
||||
use compact_space::CompactSpaceDecompressor;
|
||||
pub use monotonic_mapping::{MonotonicallyMappableToU64, StrictlyMonotonicFn};
|
||||
use monotonic_mapping::{StrictlyMonotonicMappingInverter, StrictlyMonotonicMappingToInternal};
|
||||
pub use monotonic_mapping_u128::MonotonicallyMappableToU128;
|
||||
use serialize::U128Header;
|
||||
|
||||
mod merge;
|
||||
mod compact_space;
|
||||
pub(crate) mod monotonic_mapping;
|
||||
pub(crate) mod monotonic_mapping_u128;
|
||||
mod stats;
|
||||
mod u128_based;
|
||||
mod u64_based;
|
||||
mod vec_column;
|
||||
pub(crate) mod u64_based;
|
||||
|
||||
mod monotonic_column;
|
||||
mod column;
|
||||
pub mod serialize;
|
||||
|
||||
pub(crate) use merge::MergedColumnValues;
|
||||
pub use stats::ColumnStats;
|
||||
pub use u128_based::{open_u128_mapped, serialize_column_values_u128};
|
||||
pub use serialize::serialize_column_values_u128;
|
||||
pub use stats::Stats;
|
||||
pub use u64_based::{
|
||||
load_u64_based_column_values, serialize_and_load_u64_based_column_values,
|
||||
serialize_u64_based_column_values, CodecType, ALL_U64_CODEC_TYPES,
|
||||
};
|
||||
pub use vec_column::VecColumn;
|
||||
|
||||
pub use self::monotonic_column::monotonic_map_column;
|
||||
use crate::RowId;
|
||||
pub use self::column::{monotonic_map_column, ColumnValues, IterColumn, VecColumn};
|
||||
|
||||
/// `ColumnValues` provides access to a dense field column.
|
||||
///
|
||||
/// `Column` are just a wrapper over `ColumnValues` and a `ColumnIndex`.
|
||||
///
|
||||
/// Any methods with a default and specialized implementation need to be called in the
|
||||
/// wrappers that implement the trait: Arc and MonotonicMappingColumn
|
||||
pub trait ColumnValues<T: PartialOrd = u64>: Send + Sync {
|
||||
/// Return the value associated with the given idx.
|
||||
///
|
||||
/// This accessor should return as fast as possible.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// May panic if `idx` is greater than the column length.
|
||||
fn get_val(&self, idx: u32) -> T;
|
||||
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
|
||||
#[repr(u8)]
|
||||
/// Available codecs to use to encode the u128 (via [`MonotonicallyMappableToU128`]) converted data.
|
||||
pub enum U128FastFieldCodecType {
|
||||
/// This codec takes a large number space (u128) and reduces it to a compact number space, by
|
||||
/// removing the holes.
|
||||
CompactSpace = 1,
|
||||
}
|
||||
|
||||
/// Allows to push down multiple fetch calls, to avoid dynamic dispatch overhead.
|
||||
///
|
||||
/// idx and output should have the same length
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// May panic if `idx` is greater than the column length.
|
||||
fn get_vals(&self, indexes: &[u32], output: &mut [T]) {
|
||||
assert!(indexes.len() == output.len());
|
||||
let out_and_idx_chunks = output.chunks_exact_mut(4).zip(indexes.chunks_exact(4));
|
||||
for (out_x4, idx_x4) in out_and_idx_chunks {
|
||||
out_x4[0] = self.get_val(idx_x4[0]);
|
||||
out_x4[1] = self.get_val(idx_x4[1]);
|
||||
out_x4[2] = self.get_val(idx_x4[2]);
|
||||
out_x4[3] = self.get_val(idx_x4[3]);
|
||||
}
|
||||
|
||||
let step_size = 4;
|
||||
let cutoff = indexes.len() - indexes.len() % step_size;
|
||||
|
||||
for idx in cutoff..indexes.len() {
|
||||
output[idx] = self.get_val(indexes[idx]);
|
||||
}
|
||||
impl BinarySerializable for U128FastFieldCodecType {
|
||||
fn serialize<W: Write + ?Sized>(&self, wrt: &mut W) -> io::Result<()> {
|
||||
self.to_code().serialize(wrt)
|
||||
}
|
||||
|
||||
/// Fills an output buffer with the fast field values
|
||||
/// associated with the `DocId` going from
|
||||
/// `start` to `start + output.len()`.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Must panic if `start + output.len()` is greater than
|
||||
/// the segment's `maxdoc`.
|
||||
#[inline(always)]
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
for (out, idx) in output.iter_mut().zip(start..) {
|
||||
*out = self.get_val(idx as u32);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the row ids of values which are in the provided value range.
|
||||
///
|
||||
/// Note that position == docid for single value fast fields
|
||||
#[inline(always)]
|
||||
fn get_row_ids_for_value_range(
|
||||
&self,
|
||||
value_range: RangeInclusive<T>,
|
||||
row_id_range: Range<RowId>,
|
||||
row_id_hits: &mut Vec<RowId>,
|
||||
) {
|
||||
let row_id_range = row_id_range.start..row_id_range.end.min(self.num_vals());
|
||||
for idx in row_id_range.start..row_id_range.end {
|
||||
let val = self.get_val(idx);
|
||||
if value_range.contains(&val) {
|
||||
row_id_hits.push(idx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a lower bound for this column of values.
|
||||
///
|
||||
/// All values are guaranteed to be higher than `.min_value()`
|
||||
/// but this value is not necessary the best boundary value.
|
||||
///
|
||||
/// We have
|
||||
/// ∀i < self.num_vals(), self.get_val(i) >= self.min_value()
|
||||
/// But we don't have necessarily
|
||||
/// ∃i < self.num_vals(), self.get_val(i) == self.min_value()
|
||||
fn min_value(&self) -> T;
|
||||
|
||||
/// Returns an upper bound for this column of values.
|
||||
///
|
||||
/// All values are guaranteed to be lower than `.max_value()`
|
||||
/// but this value is not necessary the best boundary value.
|
||||
///
|
||||
/// We have
|
||||
/// ∀i < self.num_vals(), self.get_val(i) <= self.max_value()
|
||||
/// But we don't have necessarily
|
||||
/// ∃i < self.num_vals(), self.get_val(i) == self.max_value()
|
||||
fn max_value(&self) -> T;
|
||||
|
||||
/// The number of values in the column.
|
||||
fn num_vals(&self) -> u32;
|
||||
|
||||
/// Returns a iterator over the data
|
||||
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = T> + 'a> {
|
||||
Box::new((0..self.num_vals()).map(|idx| self.get_val(idx)))
|
||||
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
|
||||
let code = u8::deserialize(reader)?;
|
||||
let codec_type: Self = Self::from_code(code)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Unknown code `{code}.`"))?;
|
||||
Ok(codec_type)
|
||||
}
|
||||
}
|
||||
|
||||
/// Empty column of values.
|
||||
pub struct EmptyColumnValues;
|
||||
|
||||
impl<T: PartialOrd + Default> ColumnValues<T> for EmptyColumnValues {
|
||||
fn get_val(&self, _idx: u32) -> T {
|
||||
panic!("Internal Error: Called get_val of empty column.")
|
||||
impl U128FastFieldCodecType {
|
||||
pub(crate) fn to_code(self) -> u8 {
|
||||
self as u8
|
||||
}
|
||||
|
||||
fn min_value(&self) -> T {
|
||||
T::default()
|
||||
}
|
||||
|
||||
fn max_value(&self) -> T {
|
||||
T::default()
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
0
|
||||
pub(crate) fn from_code(code: u8) -> Option<Self> {
|
||||
match code {
|
||||
1 => Some(Self::CompactSpace),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Copy + PartialOrd + Debug> ColumnValues<T> for Arc<dyn ColumnValues<T>> {
|
||||
#[inline(always)]
|
||||
fn get_val(&self, idx: u32) -> T {
|
||||
self.as_ref().get_val(idx)
|
||||
}
|
||||
/// Returns the correct codec reader wrapped in the `Arc` for the data.
|
||||
pub fn open_u128_mapped<T: MonotonicallyMappableToU128 + Debug>(
|
||||
mut bytes: OwnedBytes,
|
||||
) -> io::Result<Arc<dyn ColumnValues<T>>> {
|
||||
let header = U128Header::deserialize(&mut bytes)?;
|
||||
assert_eq!(header.codec_type, U128FastFieldCodecType::CompactSpace);
|
||||
let reader = CompactSpaceDecompressor::open(bytes)?;
|
||||
|
||||
#[inline(always)]
|
||||
fn min_value(&self) -> T {
|
||||
self.as_ref().min_value()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn max_value(&self) -> T {
|
||||
self.as_ref().max_value()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.as_ref().num_vals()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn iter<'b>(&'b self) -> Box<dyn Iterator<Item = T> + 'b> {
|
||||
self.as_ref().iter()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
self.as_ref().get_range(start, output)
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn get_row_ids_for_value_range(
|
||||
&self,
|
||||
range: RangeInclusive<T>,
|
||||
doc_id_range: Range<u32>,
|
||||
positions: &mut Vec<u32>,
|
||||
) {
|
||||
self.as_ref()
|
||||
.get_row_ids_for_value_range(range, doc_id_range, positions)
|
||||
}
|
||||
let inverted: StrictlyMonotonicMappingInverter<StrictlyMonotonicMappingToInternal<T>> =
|
||||
StrictlyMonotonicMappingToInternal::<T>::new().into();
|
||||
Ok(Arc::new(monotonic_map_column(reader, inverted)))
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "unstable"))]
|
||||
mod bench;
|
||||
mod bench {
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::OwnedBytes;
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use test::{self, Bencher};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn get_data() -> Vec<u64> {
|
||||
let mut rng = StdRng::seed_from_u64(2u64);
|
||||
let mut data: Vec<_> = (100..55000_u64)
|
||||
.map(|num| num + rng.gen::<u8>() as u64)
|
||||
.collect();
|
||||
data.push(99_000);
|
||||
data.insert(1000, 2000);
|
||||
data.insert(2000, 100);
|
||||
data.insert(3000, 4100);
|
||||
data.insert(4000, 100);
|
||||
data.insert(5000, 800);
|
||||
data
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn value_iter() -> impl Iterator<Item = u64> {
|
||||
0..20_000
|
||||
}
|
||||
fn get_reader_for_bench<Codec: FastFieldCodec>(data: &[u64]) -> Codec::Reader {
|
||||
let mut bytes = Vec::new();
|
||||
let min_value = *data.iter().min().unwrap();
|
||||
let data = data.iter().map(|el| *el - min_value).collect::<Vec<_>>();
|
||||
let col = VecColumn::from(&data);
|
||||
let normalized_header = NormalizedHeader {
|
||||
num_vals: col.num_vals(),
|
||||
max_value: col.max_value(),
|
||||
};
|
||||
Codec::serialize(&VecColumn::from(&data), &mut bytes).unwrap();
|
||||
Codec::open_from_bytes(OwnedBytes::new(bytes), normalized_header).unwrap()
|
||||
}
|
||||
fn bench_get<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
|
||||
let col = get_reader_for_bench::<Codec>(data);
|
||||
b.iter(|| {
|
||||
let mut sum = 0u64;
|
||||
for pos in value_iter() {
|
||||
let val = col.get_val(pos as u32);
|
||||
sum = sum.wrapping_add(val);
|
||||
}
|
||||
sum
|
||||
});
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn bench_get_dynamic_helper(b: &mut Bencher, col: Arc<dyn ColumnValues>) {
|
||||
b.iter(|| {
|
||||
let mut sum = 0u64;
|
||||
for pos in value_iter() {
|
||||
let val = col.get_val(pos as u32);
|
||||
sum = sum.wrapping_add(val);
|
||||
}
|
||||
sum
|
||||
});
|
||||
}
|
||||
|
||||
fn bench_get_dynamic<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
|
||||
let col = Arc::new(get_reader_for_bench::<Codec>(data));
|
||||
bench_get_dynamic_helper(b, col);
|
||||
}
|
||||
fn bench_create<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
|
||||
let min_value = *data.iter().min().unwrap();
|
||||
let data = data.iter().map(|el| *el - min_value).collect::<Vec<_>>();
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
b.iter(|| {
|
||||
bytes.clear();
|
||||
Codec::serialize(&VecColumn::from(&data), &mut bytes).unwrap();
|
||||
});
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_fastfield_bitpack_create(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_create::<BitpackedCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_linearinterpol_create(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_create::<LinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_multilinearinterpol_create(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_create::<BlockwiseLinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_bitpack_get(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get::<BitpackedCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_bitpack_get_dynamic(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get_dynamic::<BitpackedCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_linearinterpol_get(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get::<LinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_linearinterpol_get_dynamic(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get_dynamic::<LinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_multilinearinterpol_get(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get::<BlockwiseLinearCodec>(b, &data);
|
||||
}
|
||||
#[bench]
|
||||
fn bench_fastfield_multilinearinterpol_get_dynamic(b: &mut Bencher) {
|
||||
let data: Vec<_> = get_data();
|
||||
bench_get_dynamic::<BlockwiseLinearCodec>(b, &data);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
use std::fmt::Debug;
|
||||
use std::marker::PhantomData;
|
||||
use std::ops::{Range, RangeInclusive};
|
||||
|
||||
use crate::column_values::monotonic_mapping::StrictlyMonotonicFn;
|
||||
use crate::ColumnValues;
|
||||
|
||||
struct MonotonicMappingColumn<C, T, Input> {
|
||||
from_column: C,
|
||||
monotonic_mapping: T,
|
||||
_phantom: PhantomData<Input>,
|
||||
}
|
||||
|
||||
/// Creates a view of a column transformed by a strictly monotonic mapping. See
|
||||
/// [`StrictlyMonotonicFn`].
|
||||
///
|
||||
/// E.g. apply a gcd monotonic_mapping([100, 200, 300]) == [1, 2, 3]
|
||||
/// monotonic_mapping.mapping() is expected to be injective, and we should always have
|
||||
/// monotonic_mapping.inverse(monotonic_mapping.mapping(el)) == el
|
||||
///
|
||||
/// The inverse of the mapping is required for:
|
||||
/// `fn get_positions_for_value_range(&self, range: RangeInclusive<T>) -> Vec<u64> `
|
||||
/// The user provides the original value range and we need to monotonic map them in the same way the
|
||||
/// serialization does before calling the underlying column.
|
||||
///
|
||||
/// Note that when opening a codec, the monotonic_mapping should be the inverse of the mapping
|
||||
/// during serialization. And therefore the monotonic_mapping_inv when opening is the same as
|
||||
/// monotonic_mapping during serialization.
|
||||
pub fn monotonic_map_column<C, T, Input, Output>(
|
||||
from_column: C,
|
||||
monotonic_mapping: T,
|
||||
) -> impl ColumnValues<Output>
|
||||
where
|
||||
C: ColumnValues<Input>,
|
||||
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
|
||||
Input: PartialOrd + Debug + Send + Sync + Clone,
|
||||
Output: PartialOrd + Debug + Send + Sync + Clone,
|
||||
{
|
||||
MonotonicMappingColumn {
|
||||
from_column,
|
||||
monotonic_mapping,
|
||||
_phantom: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
impl<C, T, Input, Output> ColumnValues<Output> for MonotonicMappingColumn<C, T, Input>
|
||||
where
|
||||
C: ColumnValues<Input>,
|
||||
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
|
||||
Input: PartialOrd + Send + Debug + Sync + Clone,
|
||||
Output: PartialOrd + Send + Debug + Sync + Clone,
|
||||
{
|
||||
#[inline(always)]
|
||||
fn get_val(&self, idx: u32) -> Output {
|
||||
let from_val = self.from_column.get_val(idx);
|
||||
self.monotonic_mapping.mapping(from_val)
|
||||
}
|
||||
|
||||
fn min_value(&self) -> Output {
|
||||
let from_min_value = self.from_column.min_value();
|
||||
self.monotonic_mapping.mapping(from_min_value)
|
||||
}
|
||||
|
||||
fn max_value(&self) -> Output {
|
||||
let from_max_value = self.from_column.max_value();
|
||||
self.monotonic_mapping.mapping(from_max_value)
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.from_column.num_vals()
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = Output> + '_> {
|
||||
Box::new(
|
||||
self.from_column
|
||||
.iter()
|
||||
.map(|el| self.monotonic_mapping.mapping(el)),
|
||||
)
|
||||
}
|
||||
|
||||
fn get_row_ids_for_value_range(
|
||||
&self,
|
||||
range: RangeInclusive<Output>,
|
||||
doc_id_range: Range<u32>,
|
||||
positions: &mut Vec<u32>,
|
||||
) {
|
||||
self.from_column.get_row_ids_for_value_range(
|
||||
self.monotonic_mapping.inverse(range.start().clone())
|
||||
..=self.monotonic_mapping.inverse(range.end().clone()),
|
||||
doc_id_range,
|
||||
positions,
|
||||
)
|
||||
}
|
||||
|
||||
// We voluntarily do not implement get_range as it yields a regression,
|
||||
// and we do not have any specialized implementation anyway.
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::column_values::monotonic_mapping::{
|
||||
StrictlyMonotonicMappingInverter, StrictlyMonotonicMappingToInternal,
|
||||
};
|
||||
use crate::column_values::VecColumn;
|
||||
|
||||
#[test]
|
||||
fn test_monotonic_mapping_iter() {
|
||||
let vals: Vec<u64> = (0..100u64).map(|el| el * 10).collect();
|
||||
let col = VecColumn::from(&vals);
|
||||
let mapped = monotonic_map_column(
|
||||
col,
|
||||
StrictlyMonotonicMappingInverter::from(StrictlyMonotonicMappingToInternal::<i64>::new()),
|
||||
);
|
||||
let val_i64s: Vec<u64> = mapped.iter().collect();
|
||||
for i in 0..100 {
|
||||
assert_eq!(val_i64s[i as usize], mapped.get_val(i));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
use std::fmt::Debug;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use common::DateTime;
|
||||
use fastdivide::DividerU64;
|
||||
|
||||
use super::MonotonicallyMappableToU128;
|
||||
use crate::RowId;
|
||||
@@ -112,6 +112,65 @@ where T: MonotonicallyMappableToU64
|
||||
}
|
||||
}
|
||||
|
||||
/// Mapping dividing by gcd and a base value.
|
||||
///
|
||||
/// The function is assumed to be only called on values divided by passed
|
||||
/// gcd value. (It is necessary for the function to be monotonic.)
|
||||
pub(crate) struct StrictlyMonotonicMappingToInternalGCDBaseval {
|
||||
gcd_divider: DividerU64,
|
||||
gcd: u64,
|
||||
min_value: u64,
|
||||
}
|
||||
impl StrictlyMonotonicMappingToInternalGCDBaseval {
|
||||
pub(crate) fn new(gcd: u64, min_value: u64) -> Self {
|
||||
let gcd_divider = DividerU64::divide_by(gcd);
|
||||
Self {
|
||||
gcd_divider,
|
||||
gcd,
|
||||
min_value,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
|
||||
for StrictlyMonotonicMappingToInternalGCDBaseval
|
||||
{
|
||||
#[inline(always)]
|
||||
fn mapping(&self, inp: External) -> u64 {
|
||||
self.gcd_divider
|
||||
.divide(External::to_u64(inp) - self.min_value)
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn inverse(&self, out: u64) -> External {
|
||||
External::from_u64(self.min_value + out * self.gcd)
|
||||
}
|
||||
}
|
||||
|
||||
/// Strictly monotonic mapping with a base value.
|
||||
pub(crate) struct StrictlyMonotonicMappingToInternalBaseval {
|
||||
min_value: u64,
|
||||
}
|
||||
impl StrictlyMonotonicMappingToInternalBaseval {
|
||||
#[inline(always)]
|
||||
pub(crate) fn new(min_value: u64) -> Self {
|
||||
Self { min_value }
|
||||
}
|
||||
}
|
||||
|
||||
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
|
||||
for StrictlyMonotonicMappingToInternalBaseval
|
||||
{
|
||||
#[inline(always)]
|
||||
fn mapping(&self, val: External) -> u64 {
|
||||
External::to_u64(val) - self.min_value
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn inverse(&self, val: u64) -> External {
|
||||
External::from_u64(self.min_value + val)
|
||||
}
|
||||
}
|
||||
|
||||
impl MonotonicallyMappableToU64 for u64 {
|
||||
#[inline(always)]
|
||||
fn to_u64(self) -> u64 {
|
||||
@@ -136,15 +195,17 @@ impl MonotonicallyMappableToU64 for i64 {
|
||||
}
|
||||
}
|
||||
|
||||
impl MonotonicallyMappableToU64 for DateTime {
|
||||
impl MonotonicallyMappableToU64 for crate::DateTime {
|
||||
#[inline(always)]
|
||||
fn to_u64(self) -> u64 {
|
||||
common::i64_to_u64(self.into_timestamp_micros())
|
||||
common::i64_to_u64(self.timestamp_micros)
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn from_u64(val: u64) -> Self {
|
||||
DateTime::from_timestamp_micros(common::u64_to_i64(val))
|
||||
crate::DateTime {
|
||||
timestamp_micros: common::u64_to_i64(val),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -200,6 +261,13 @@ mod tests {
|
||||
// TODO
|
||||
// identity mapping
|
||||
// test_round_trip(&StrictlyMonotonicMappingToInternal::<u128>::new(), 100u128);
|
||||
|
||||
// base value to i64 round trip
|
||||
let mapping = StrictlyMonotonicMappingToInternalBaseval::new(100);
|
||||
test_round_trip::<_, _, u64>(&mapping, 100i64);
|
||||
// base value and gcd to u64 round trip
|
||||
let mapping = StrictlyMonotonicMappingToInternalGCDBaseval::new(10, 100);
|
||||
test_round_trip::<_, _, u64>(&mapping, 100u64);
|
||||
}
|
||||
|
||||
fn test_round_trip<T: StrictlyMonotonicFn<K, L>, K: std::fmt::Debug + Eq + Copy, L>(
|
||||
|
||||
@@ -1,19 +1,25 @@
|
||||
use std::fmt::Debug;
|
||||
use std::io;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
|
||||
mod compact_space;
|
||||
use common::{BinarySerializable, VInt};
|
||||
|
||||
use common::{BinarySerializable, OwnedBytes, VInt};
|
||||
use compact_space::{CompactSpaceCompressor, CompactSpaceDecompressor};
|
||||
|
||||
use crate::column_values::monotonic_map_column;
|
||||
use crate::column_values::monotonic_mapping::{
|
||||
StrictlyMonotonicMappingInverter, StrictlyMonotonicMappingToInternal,
|
||||
};
|
||||
use crate::column_values::compact_space::CompactSpaceCompressor;
|
||||
use crate::column_values::U128FastFieldCodecType;
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{ColumnValues, MonotonicallyMappableToU128};
|
||||
use crate::MonotonicallyMappableToU128;
|
||||
|
||||
/// The normalized header gives some parameters after applying the following
|
||||
/// normalization of the vector:
|
||||
/// `val -> (val - min_value) / gcd`
|
||||
///
|
||||
/// By design, after normalization, `min_value = 0` and `gcd = 1`.
|
||||
#[derive(Debug, Copy, Clone)]
|
||||
pub struct NormalizedHeader {
|
||||
/// The number of values in the underlying column.
|
||||
pub num_vals: u32,
|
||||
/// The max value of the underlying column.
|
||||
pub max_value: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct U128Header {
|
||||
@@ -41,18 +47,20 @@ impl BinarySerializable for U128Header {
|
||||
/// Serializes u128 values with the compact space codec.
|
||||
pub fn serialize_column_values_u128<T: MonotonicallyMappableToU128>(
|
||||
iterable: &dyn Iterable<T>,
|
||||
num_vals: u32,
|
||||
output: &mut impl io::Write,
|
||||
) -> io::Result<()> {
|
||||
let header = U128Header {
|
||||
num_vals,
|
||||
codec_type: U128FastFieldCodecType::CompactSpace,
|
||||
};
|
||||
header.serialize(output)?;
|
||||
let compressor = CompactSpaceCompressor::train_from(
|
||||
iterable
|
||||
.boxed_iter()
|
||||
.map(MonotonicallyMappableToU128::to_u128),
|
||||
num_vals,
|
||||
);
|
||||
let header = U128Header {
|
||||
num_vals: compressor.num_vals(),
|
||||
codec_type: U128FastFieldCodecType::CompactSpace,
|
||||
};
|
||||
header.serialize(output)?;
|
||||
compressor.compress_into(
|
||||
iterable
|
||||
.boxed_iter()
|
||||
@@ -62,57 +70,11 @@ pub fn serialize_column_values_u128<T: MonotonicallyMappableToU128>(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
|
||||
#[repr(u8)]
|
||||
/// Available codecs to use to encode the u128 (via [`MonotonicallyMappableToU128`]) converted data.
|
||||
pub(crate) enum U128FastFieldCodecType {
|
||||
/// This codec takes a large number space (u128) and reduces it to a compact number space, by
|
||||
/// removing the holes.
|
||||
CompactSpace = 1,
|
||||
}
|
||||
|
||||
impl BinarySerializable for U128FastFieldCodecType {
|
||||
fn serialize<W: Write + ?Sized>(&self, wrt: &mut W) -> io::Result<()> {
|
||||
self.to_code().serialize(wrt)
|
||||
}
|
||||
|
||||
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
|
||||
let code = u8::deserialize(reader)?;
|
||||
let codec_type: Self = Self::from_code(code)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Unknown code `{code}.`"))?;
|
||||
Ok(codec_type)
|
||||
}
|
||||
}
|
||||
|
||||
impl U128FastFieldCodecType {
|
||||
pub(crate) fn to_code(self) -> u8 {
|
||||
self as u8
|
||||
}
|
||||
|
||||
pub(crate) fn from_code(code: u8) -> Option<Self> {
|
||||
match code {
|
||||
1 => Some(Self::CompactSpace),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the correct codec reader wrapped in the `Arc` for the data.
|
||||
pub fn open_u128_mapped<T: MonotonicallyMappableToU128 + Debug>(
|
||||
mut bytes: OwnedBytes,
|
||||
) -> io::Result<Arc<dyn ColumnValues<T>>> {
|
||||
let header = U128Header::deserialize(&mut bytes)?;
|
||||
assert_eq!(header.codec_type, U128FastFieldCodecType::CompactSpace);
|
||||
let reader = CompactSpaceDecompressor::open(bytes)?;
|
||||
let inverted: StrictlyMonotonicMappingInverter<StrictlyMonotonicMappingToInternal<T>> =
|
||||
StrictlyMonotonicMappingToInternal::<T>::new().into();
|
||||
Ok(Arc::new(monotonic_map_column(reader, inverted)))
|
||||
}
|
||||
#[cfg(test)]
|
||||
pub mod tests {
|
||||
use super::*;
|
||||
use crate::column_values::u64_based::{
|
||||
serialize_and_load_u64_based_column_values, serialize_u64_based_column_values,
|
||||
self, serialize_and_load_u64_based_column_values, serialize_u64_based_column_values,
|
||||
ALL_U64_CODEC_TYPES,
|
||||
};
|
||||
use crate::column_values::CodecType;
|
||||
@@ -6,28 +6,21 @@ use common::{BinarySerializable, VInt};
|
||||
|
||||
use crate::RowId;
|
||||
|
||||
/// Column statistics.
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub struct ColumnStats {
|
||||
/// GCD of the elements `el - min(column)`.
|
||||
pub struct Stats {
|
||||
pub gcd: NonZeroU64,
|
||||
/// Minimum value of the column.
|
||||
pub min_value: u64,
|
||||
/// Maximum value of the column.
|
||||
pub max_value: u64,
|
||||
/// Number of rows in the column.
|
||||
pub num_rows: RowId,
|
||||
}
|
||||
|
||||
impl ColumnStats {
|
||||
/// Amplitude of value.
|
||||
/// Difference between the maximum and the minimum value.
|
||||
impl Stats {
|
||||
pub fn amplitude(&self) -> u64 {
|
||||
self.max_value - self.min_value
|
||||
}
|
||||
}
|
||||
|
||||
impl BinarySerializable for ColumnStats {
|
||||
impl BinarySerializable for Stats {
|
||||
fn serialize<W: Write + ?Sized>(&self, writer: &mut W) -> io::Result<()> {
|
||||
VInt(self.min_value).serialize(writer)?;
|
||||
VInt(self.gcd.get()).serialize(writer)?;
|
||||
@@ -44,7 +37,7 @@ impl BinarySerializable for ColumnStats {
|
||||
let amplitude = VInt::deserialize(reader)?.0 * gcd.get();
|
||||
let max_value = min_value + amplitude;
|
||||
let num_rows = VInt::deserialize(reader)?.0 as RowId;
|
||||
Ok(ColumnStats {
|
||||
Ok(Stats {
|
||||
min_value,
|
||||
max_value,
|
||||
num_rows,
|
||||
@@ -59,21 +52,21 @@ mod tests {
|
||||
|
||||
use common::BinarySerializable;
|
||||
|
||||
use crate::column_values::ColumnStats;
|
||||
use crate::column_values::Stats;
|
||||
|
||||
#[track_caller]
|
||||
fn test_stats_ser_deser_aux(stats: &ColumnStats, num_bytes: usize) {
|
||||
fn test_stats_ser_deser_aux(stats: &Stats, num_bytes: usize) {
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
stats.serialize(&mut buffer).unwrap();
|
||||
assert_eq!(buffer.len(), num_bytes);
|
||||
let deser_stats = ColumnStats::deserialize(&mut &buffer[..]).unwrap();
|
||||
let deser_stats = Stats::deserialize(&mut &buffer[..]).unwrap();
|
||||
assert_eq!(stats, &deser_stats);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stats_serialization() {
|
||||
test_stats_ser_deser_aux(
|
||||
&(ColumnStats {
|
||||
&(Stats {
|
||||
gcd: NonZeroU64::new(3).unwrap(),
|
||||
min_value: 1,
|
||||
max_value: 3001,
|
||||
@@ -82,7 +75,7 @@ mod tests {
|
||||
5,
|
||||
);
|
||||
test_stats_ser_deser_aux(
|
||||
&(ColumnStats {
|
||||
&(Stats {
|
||||
gcd: NonZeroU64::new(1_000).unwrap(),
|
||||
min_value: 1,
|
||||
max_value: 3001,
|
||||
@@ -91,7 +84,7 @@ mod tests {
|
||||
5,
|
||||
);
|
||||
test_stats_ser_deser_aux(
|
||||
&(ColumnStats {
|
||||
&(Stats {
|
||||
gcd: NonZeroU64::new(1).unwrap(),
|
||||
min_value: 0,
|
||||
max_value: 0,
|
||||
|
||||
@@ -4,7 +4,7 @@ use common::{BinarySerializable, OwnedBytes};
|
||||
use fastdivide::DividerU64;
|
||||
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
|
||||
|
||||
use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, ColumnStats};
|
||||
use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, Stats};
|
||||
use crate::{ColumnValues, RowId};
|
||||
|
||||
/// Depending on the field type, a different
|
||||
@@ -13,7 +13,7 @@ use crate::{ColumnValues, RowId};
|
||||
pub struct BitpackedReader {
|
||||
data: OwnedBytes,
|
||||
bit_unpacker: BitUnpacker,
|
||||
stats: ColumnStats,
|
||||
stats: Stats,
|
||||
}
|
||||
|
||||
impl ColumnValues for BitpackedReader {
|
||||
@@ -36,7 +36,7 @@ impl ColumnValues for BitpackedReader {
|
||||
}
|
||||
}
|
||||
|
||||
fn num_bits(stats: &ColumnStats) -> u8 {
|
||||
fn num_bits(stats: &Stats) -> u8 {
|
||||
compute_num_bits(stats.amplitude() / stats.gcd)
|
||||
}
|
||||
|
||||
@@ -46,14 +46,14 @@ pub struct BitpackedCodecEstimator;
|
||||
impl ColumnCodecEstimator for BitpackedCodecEstimator {
|
||||
fn collect(&mut self, _value: u64) {}
|
||||
|
||||
fn estimate(&self, stats: &ColumnStats) -> Option<u64> {
|
||||
fn estimate(&self, stats: &Stats) -> Option<u64> {
|
||||
let num_bits_per_value = num_bits(stats);
|
||||
Some(stats.num_bytes() + (stats.num_rows as u64 * (num_bits_per_value as u64) + 7) / 8)
|
||||
}
|
||||
|
||||
fn serialize(
|
||||
&self,
|
||||
stats: &ColumnStats,
|
||||
stats: &Stats,
|
||||
vals: &mut dyn Iterator<Item = u64>,
|
||||
wrt: &mut dyn Write,
|
||||
) -> io::Result<()> {
|
||||
@@ -72,12 +72,12 @@ impl ColumnCodecEstimator for BitpackedCodecEstimator {
|
||||
pub struct BitpackedCodec;
|
||||
|
||||
impl ColumnCodec for BitpackedCodec {
|
||||
type ColumnValues = BitpackedReader;
|
||||
type Reader = BitpackedReader;
|
||||
type Estimator = BitpackedCodecEstimator;
|
||||
|
||||
/// Opens a fast field given a file.
|
||||
fn load(mut data: OwnedBytes) -> io::Result<Self::ColumnValues> {
|
||||
let stats = ColumnStats::deserialize(&mut data)?;
|
||||
fn load(mut data: OwnedBytes) -> io::Result<Self::Reader> {
|
||||
let stats = Stats::deserialize(&mut data)?;
|
||||
let num_bits = num_bits(&stats);
|
||||
let bit_unpacker = BitUnpacker::new(num_bits);
|
||||
Ok(BitpackedReader {
|
||||
|
||||
@@ -7,7 +7,7 @@ use fastdivide::DividerU64;
|
||||
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
|
||||
|
||||
use crate::column_values::u64_based::line::Line;
|
||||
use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, ColumnStats};
|
||||
use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, Stats};
|
||||
use crate::column_values::{ColumnValues, VecColumn};
|
||||
use crate::MonotonicallyMappableToU64;
|
||||
|
||||
@@ -84,7 +84,7 @@ impl ColumnCodecEstimator for BlockwiseLinearEstimator {
|
||||
self.block.clear();
|
||||
}
|
||||
}
|
||||
fn estimate(&self, stats: &ColumnStats) -> Option<u64> {
|
||||
fn estimate(&self, stats: &Stats) -> Option<u64> {
|
||||
let mut estimate = 4 + stats.num_bytes() + self.meta_num_bytes + self.values_num_bytes;
|
||||
if stats.gcd.get() > 1 {
|
||||
let estimate_gain_from_gcd =
|
||||
@@ -100,7 +100,7 @@ impl ColumnCodecEstimator for BlockwiseLinearEstimator {
|
||||
|
||||
fn serialize(
|
||||
&self,
|
||||
stats: &ColumnStats,
|
||||
stats: &Stats,
|
||||
mut vals: &mut dyn Iterator<Item = u64>,
|
||||
wrt: &mut dyn Write,
|
||||
) -> io::Result<()> {
|
||||
@@ -165,12 +165,12 @@ impl ColumnCodecEstimator for BlockwiseLinearEstimator {
|
||||
pub struct BlockwiseLinearCodec;
|
||||
|
||||
impl ColumnCodec<u64> for BlockwiseLinearCodec {
|
||||
type ColumnValues = BlockwiseLinearReader;
|
||||
type Reader = BlockwiseLinearReader;
|
||||
|
||||
type Estimator = BlockwiseLinearEstimator;
|
||||
|
||||
fn load(mut bytes: OwnedBytes) -> io::Result<Self::ColumnValues> {
|
||||
let stats = ColumnStats::deserialize(&mut bytes)?;
|
||||
fn load(mut bytes: OwnedBytes) -> io::Result<Self::Reader> {
|
||||
let stats = Stats::deserialize(&mut bytes)?;
|
||||
let footer_len: u32 = (&bytes[bytes.len() - 4..]).deserialize()?;
|
||||
let footer_offset = bytes.len() - 4 - footer_len as usize;
|
||||
let (data, mut footer) = bytes.split(footer_offset);
|
||||
@@ -195,14 +195,14 @@ impl ColumnCodec<u64> for BlockwiseLinearCodec {
|
||||
pub struct BlockwiseLinearReader {
|
||||
blocks: Arc<[Block]>,
|
||||
data: OwnedBytes,
|
||||
stats: ColumnStats,
|
||||
stats: Stats,
|
||||
}
|
||||
|
||||
impl ColumnValues for BlockwiseLinearReader {
|
||||
#[inline(always)]
|
||||
fn get_val(&self, idx: u32) -> u64 {
|
||||
let block_id = (idx / BLOCK_SIZE) as usize;
|
||||
let idx_within_block = idx % BLOCK_SIZE;
|
||||
let block_id = (idx / BLOCK_SIZE as u32) as usize;
|
||||
let idx_within_block = idx % (BLOCK_SIZE as u32);
|
||||
let block = &self.blocks[block_id];
|
||||
let interpoled_val: u64 = block.line.eval(idx_within_block);
|
||||
let block_bytes = &self.data[block.data_start_offset..];
|
||||
|
||||
@@ -67,6 +67,19 @@ impl Line {
|
||||
self.intercept.wrapping_add(linear_part)
|
||||
}
|
||||
|
||||
// Same as train, but the intercept is only estimated from provided sample positions
|
||||
pub fn estimate(sample_positions_and_values: &[(u64, u64)]) -> Self {
|
||||
let first_val = sample_positions_and_values[0].1;
|
||||
let last_val = sample_positions_and_values[sample_positions_and_values.len() - 1].1;
|
||||
let num_vals = sample_positions_and_values[sample_positions_and_values.len() - 1].0 + 1;
|
||||
Self::train_from(
|
||||
first_val,
|
||||
last_val,
|
||||
num_vals as u32,
|
||||
sample_positions_and_values.iter().cloned(),
|
||||
)
|
||||
}
|
||||
|
||||
// Intercept is only computed from provided positions
|
||||
pub fn train_from(
|
||||
first_val: u64,
|
||||
|
||||
@@ -5,9 +5,9 @@ use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
|
||||
|
||||
use super::line::Line;
|
||||
use super::ColumnValues;
|
||||
use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, ColumnStats};
|
||||
use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, Stats};
|
||||
use crate::column_values::VecColumn;
|
||||
use crate::RowId;
|
||||
use crate::{MonotonicallyMappableToU64, RowId};
|
||||
|
||||
const HALF_SPACE: u64 = u64::MAX / 2;
|
||||
const LINE_ESTIMATION_BLOCK_LEN: usize = 512;
|
||||
@@ -18,7 +18,7 @@ const LINE_ESTIMATION_BLOCK_LEN: usize = 512;
|
||||
pub struct LinearReader {
|
||||
data: OwnedBytes,
|
||||
linear_params: LinearParams,
|
||||
stats: ColumnStats,
|
||||
stats: Stats,
|
||||
}
|
||||
|
||||
impl ColumnValues for LinearReader {
|
||||
@@ -106,7 +106,7 @@ impl ColumnCodecEstimator for LinearCodecEstimator {
|
||||
}
|
||||
}
|
||||
|
||||
fn estimate(&self, stats: &ColumnStats) -> Option<u64> {
|
||||
fn estimate(&self, stats: &Stats) -> Option<u64> {
|
||||
let line = self.line?;
|
||||
let amplitude = self.max_deviation - self.min_deviation;
|
||||
let num_bits = compute_num_bits(amplitude);
|
||||
@@ -123,7 +123,7 @@ impl ColumnCodecEstimator for LinearCodecEstimator {
|
||||
|
||||
fn serialize(
|
||||
&self,
|
||||
stats: &ColumnStats,
|
||||
stats: &Stats,
|
||||
vals: &mut dyn Iterator<Item = u64>,
|
||||
wrt: &mut dyn io::Write,
|
||||
) -> io::Result<()> {
|
||||
@@ -184,12 +184,12 @@ impl LinearCodecEstimator {
|
||||
}
|
||||
|
||||
impl ColumnCodec for LinearCodec {
|
||||
type ColumnValues = LinearReader;
|
||||
type Reader = LinearReader;
|
||||
|
||||
type Estimator = LinearCodecEstimator;
|
||||
|
||||
fn load(mut data: OwnedBytes) -> io::Result<Self::ColumnValues> {
|
||||
let stats = ColumnStats::deserialize(&mut data)?;
|
||||
fn load(mut data: OwnedBytes) -> io::Result<Self::Reader> {
|
||||
let stats = Stats::deserialize(&mut data)?;
|
||||
let linear_params = LinearParams::deserialize(&mut data)?;
|
||||
Ok(LinearReader {
|
||||
stats,
|
||||
|
||||
@@ -13,61 +13,35 @@ use common::{BinarySerializable, OwnedBytes};
|
||||
use crate::column_values::monotonic_mapping::{
|
||||
StrictlyMonotonicMappingInverter, StrictlyMonotonicMappingToInternal,
|
||||
};
|
||||
pub use crate::column_values::u64_based::bitpacked::BitpackedCodec;
|
||||
pub use crate::column_values::u64_based::blockwise_linear::BlockwiseLinearCodec;
|
||||
pub use crate::column_values::u64_based::linear::LinearCodec;
|
||||
pub use crate::column_values::u64_based::stats_collector::StatsCollector;
|
||||
use crate::column_values::{monotonic_map_column, ColumnStats};
|
||||
use crate::column_values::u64_based::bitpacked::BitpackedCodec;
|
||||
use crate::column_values::u64_based::blockwise_linear::BlockwiseLinearCodec;
|
||||
use crate::column_values::u64_based::linear::LinearCodec;
|
||||
use crate::column_values::u64_based::stats_collector::StatsCollector;
|
||||
use crate::column_values::{monotonic_map_column, Stats};
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{ColumnValues, MonotonicallyMappableToU64};
|
||||
|
||||
/// A `ColumnCodecEstimator` is in charge of gathering all
|
||||
/// data required to serialize a column.
|
||||
///
|
||||
/// This happens during a first pass on data of the column elements.
|
||||
/// During that pass, all column estimators receive a call to their
|
||||
/// `.collect(el)`.
|
||||
///
|
||||
/// After this first pass, finalize is called.
|
||||
/// `.estimate(..)` then should return an accurate estimation of the
|
||||
/// size of the serialized column (were we to pick this codec.).
|
||||
/// `.serialize(..)` then serializes the column using this codec.
|
||||
pub trait ColumnCodecEstimator<T = u64>: 'static {
|
||||
/// Records a new value for estimation.
|
||||
/// This method will be called for each element of the column during
|
||||
/// `estimation`.
|
||||
fn collect(&mut self, value: u64);
|
||||
/// Finalizes the first pass phase.
|
||||
fn estimate(&self, stats: &Stats) -> Option<u64>;
|
||||
fn finalize(&mut self) {}
|
||||
/// Returns an accurate estimation of the number of bytes that will
|
||||
/// be used to represent this column.
|
||||
fn estimate(&self, stats: &ColumnStats) -> Option<u64>;
|
||||
/// Serializes the column using the given codec.
|
||||
/// This constitutes a second pass over the columns values.
|
||||
fn serialize(
|
||||
&self,
|
||||
stats: &ColumnStats,
|
||||
stats: &Stats,
|
||||
vals: &mut dyn Iterator<Item = T>,
|
||||
wrt: &mut dyn io::Write,
|
||||
) -> io::Result<()>;
|
||||
}
|
||||
|
||||
/// A column codec describes a colunm serialization format.
|
||||
pub trait ColumnCodec<T: PartialOrd = u64> {
|
||||
/// Specialized `ColumnValues` type.
|
||||
type ColumnValues: ColumnValues<T> + 'static;
|
||||
/// `Estimator` for the given codec.
|
||||
type Reader: ColumnValues<T> + 'static;
|
||||
type Estimator: ColumnCodecEstimator + Default;
|
||||
|
||||
/// Loads a column that has been serialized using this codec.
|
||||
fn load(bytes: OwnedBytes) -> io::Result<Self::ColumnValues>;
|
||||
fn load(bytes: OwnedBytes) -> io::Result<Self::Reader>;
|
||||
|
||||
/// Returns an estimator.
|
||||
fn estimator() -> Self::Estimator {
|
||||
Self::Estimator::default()
|
||||
}
|
||||
|
||||
/// Returns a boxed estimator.
|
||||
fn boxed_estimator() -> Box<dyn ColumnCodecEstimator> {
|
||||
Box::new(Self::estimator())
|
||||
}
|
||||
@@ -88,7 +62,6 @@ pub enum CodecType {
|
||||
BlockwiseLinear = 2u8,
|
||||
}
|
||||
|
||||
/// List of all available u64-base codecs.
|
||||
pub const ALL_U64_CODEC_TYPES: [CodecType; 3] = [
|
||||
CodecType::Bitpacked,
|
||||
CodecType::Linear,
|
||||
@@ -133,7 +106,6 @@ fn load_specific_codec<C: ColumnCodec, T: MonotonicallyMappableToU64>(
|
||||
}
|
||||
|
||||
impl CodecType {
|
||||
/// Returns a boxed codec estimator associated to a given `CodecType`.
|
||||
pub fn estimator(&self) -> Box<dyn ColumnCodecEstimator> {
|
||||
match self {
|
||||
CodecType::Bitpacked => BitpackedCodec::boxed_estimator(),
|
||||
@@ -143,8 +115,7 @@ impl CodecType {
|
||||
}
|
||||
}
|
||||
|
||||
/// Serializes a given column of u64-mapped values.
|
||||
pub fn serialize_u64_based_column_values<T: MonotonicallyMappableToU64>(
|
||||
pub fn serialize_u64_based_column_values<'a, T: MonotonicallyMappableToU64>(
|
||||
vals: &dyn Iterable<T>,
|
||||
codec_types: &[CodecType],
|
||||
wrt: &mut dyn Write,
|
||||
@@ -185,14 +156,11 @@ pub fn serialize_u64_based_column_values<T: MonotonicallyMappableToU64>(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load u64-based column values.
|
||||
///
|
||||
/// This method first identifies the codec off the first byte.
|
||||
pub fn load_u64_based_column_values<T: MonotonicallyMappableToU64>(
|
||||
mut bytes: OwnedBytes,
|
||||
) -> io::Result<Arc<dyn ColumnValues<T>>> {
|
||||
let codec_type: CodecType = bytes
|
||||
.first()
|
||||
.get(0)
|
||||
.copied()
|
||||
.and_then(CodecType::try_from_code)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Failed to read codec type"))?;
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::num::NonZeroU64;
|
||||
|
||||
use fastdivide::DividerU64;
|
||||
|
||||
use crate::column_values::ColumnStats;
|
||||
use crate::column_values::Stats;
|
||||
use crate::RowId;
|
||||
|
||||
/// Compute the gcd of two non null numbers.
|
||||
@@ -33,14 +33,14 @@ pub struct StatsCollector {
|
||||
}
|
||||
|
||||
impl StatsCollector {
|
||||
pub fn stats(&self) -> ColumnStats {
|
||||
pub fn stats(&self) -> Stats {
|
||||
let (min_value, max_value) = self.min_max_opt.unwrap_or((0u64, 0u64));
|
||||
let increment_gcd = if let Some((increment_gcd, _)) = self.increment_gcd_opt {
|
||||
increment_gcd
|
||||
} else {
|
||||
NonZeroU64::new(1u64).unwrap()
|
||||
};
|
||||
ColumnStats {
|
||||
Stats {
|
||||
min_value,
|
||||
max_value,
|
||||
num_rows: self.num_rows,
|
||||
@@ -97,9 +97,9 @@ mod tests {
|
||||
use std::num::NonZeroU64;
|
||||
|
||||
use crate::column_values::u64_based::stats_collector::{compute_gcd, StatsCollector};
|
||||
use crate::column_values::u64_based::ColumnStats;
|
||||
use crate::column_values::u64_based::Stats;
|
||||
|
||||
fn compute_stats(vals: impl Iterator<Item = u64>) -> ColumnStats {
|
||||
fn compute_stats(vals: impl Iterator<Item = u64>) -> Stats {
|
||||
let mut stats_collector = StatsCollector::default();
|
||||
for val in vals {
|
||||
stats_collector.collect(val);
|
||||
@@ -144,7 +144,7 @@ mod tests {
|
||||
fn test_stats() {
|
||||
assert_eq!(
|
||||
compute_stats([].into_iter()),
|
||||
ColumnStats {
|
||||
Stats {
|
||||
gcd: NonZeroU64::new(1).unwrap(),
|
||||
min_value: 0,
|
||||
max_value: 0,
|
||||
@@ -153,7 +153,7 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
compute_stats([0, 1].into_iter()),
|
||||
ColumnStats {
|
||||
Stats {
|
||||
gcd: NonZeroU64::new(1).unwrap(),
|
||||
min_value: 0,
|
||||
max_value: 1,
|
||||
@@ -162,7 +162,7 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
compute_stats([0, 1].into_iter()),
|
||||
ColumnStats {
|
||||
Stats {
|
||||
gcd: NonZeroU64::new(1).unwrap(),
|
||||
min_value: 0,
|
||||
max_value: 1,
|
||||
@@ -171,7 +171,7 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
compute_stats([10, 20, 30].into_iter()),
|
||||
ColumnStats {
|
||||
Stats {
|
||||
gcd: NonZeroU64::new(10).unwrap(),
|
||||
min_value: 10,
|
||||
max_value: 30,
|
||||
@@ -180,7 +180,7 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
compute_stats([10, 50, 10, 30].into_iter()),
|
||||
ColumnStats {
|
||||
Stats {
|
||||
gcd: NonZeroU64::new(20).unwrap(),
|
||||
min_value: 10,
|
||||
max_value: 50,
|
||||
@@ -189,7 +189,7 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
compute_stats([10, 0, 30].into_iter()),
|
||||
ColumnStats {
|
||||
Stats {
|
||||
gcd: NonZeroU64::new(10).unwrap(),
|
||||
min_value: 0,
|
||||
max_value: 30,
|
||||
|
||||
@@ -19,62 +19,6 @@ fn test_serialize_and_load_simple() {
|
||||
assert_eq!(col.get_val(1), 2);
|
||||
assert_eq!(col.get_val(2), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_column_i64() {
|
||||
let vals: [i64; 0] = [];
|
||||
let mut num_acceptable_codecs = 0;
|
||||
for codec in ALL_U64_CODEC_TYPES {
|
||||
let mut buffer = Vec::new();
|
||||
if serialize_u64_based_column_values(&&vals[..], &[codec], &mut buffer).is_err() {
|
||||
continue;
|
||||
}
|
||||
num_acceptable_codecs += 1;
|
||||
let col = load_u64_based_column_values::<i64>(OwnedBytes::new(buffer)).unwrap();
|
||||
assert_eq!(col.num_vals(), 0);
|
||||
assert_eq!(col.min_value(), i64::MIN);
|
||||
assert_eq!(col.max_value(), i64::MIN);
|
||||
}
|
||||
assert!(num_acceptable_codecs > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_column_u64() {
|
||||
let vals: [u64; 0] = [];
|
||||
let mut num_acceptable_codecs = 0;
|
||||
for codec in ALL_U64_CODEC_TYPES {
|
||||
let mut buffer = Vec::new();
|
||||
if serialize_u64_based_column_values(&&vals[..], &[codec], &mut buffer).is_err() {
|
||||
continue;
|
||||
}
|
||||
num_acceptable_codecs += 1;
|
||||
let col = load_u64_based_column_values::<u64>(OwnedBytes::new(buffer)).unwrap();
|
||||
assert_eq!(col.num_vals(), 0);
|
||||
assert_eq!(col.min_value(), u64::MIN);
|
||||
assert_eq!(col.max_value(), u64::MIN);
|
||||
}
|
||||
assert!(num_acceptable_codecs > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_column_f64() {
|
||||
let vals: [f64; 0] = [];
|
||||
let mut num_acceptable_codecs = 0;
|
||||
for codec in ALL_U64_CODEC_TYPES {
|
||||
let mut buffer = Vec::new();
|
||||
if serialize_u64_based_column_values(&&vals[..], &[codec], &mut buffer).is_err() {
|
||||
continue;
|
||||
}
|
||||
num_acceptable_codecs += 1;
|
||||
let col = load_u64_based_column_values::<f64>(OwnedBytes::new(buffer)).unwrap();
|
||||
assert_eq!(col.num_vals(), 0);
|
||||
// FIXME. f64::MIN would be better!
|
||||
assert!(col.min_value().is_nan());
|
||||
assert!(col.max_value().is_nan());
|
||||
}
|
||||
assert!(num_acceptable_codecs > 0);
|
||||
}
|
||||
|
||||
pub(crate) fn create_and_validate<TColumnCodec: ColumnCodec>(
|
||||
vals: &[u64],
|
||||
name: &str,
|
||||
@@ -99,28 +43,14 @@ pub(crate) fn create_and_validate<TColumnCodec: ColumnCodec>(
|
||||
|
||||
let reader = TColumnCodec::load(OwnedBytes::new(buffer)).unwrap();
|
||||
assert_eq!(reader.num_vals(), vals.len() as u32);
|
||||
let mut buffer = Vec::new();
|
||||
for (doc, orig_val) in vals.iter().copied().enumerate() {
|
||||
let val = reader.get_val(doc as u32);
|
||||
assert_eq!(
|
||||
val, orig_val,
|
||||
"val `{val}` does not match orig_val {orig_val:?}, in data set {name}, data `{vals:?}`",
|
||||
);
|
||||
|
||||
buffer.resize(1, 0);
|
||||
reader.get_vals(&[doc as u32], &mut buffer);
|
||||
let val = buffer[0];
|
||||
assert_eq!(
|
||||
val, orig_val,
|
||||
"val `{val}` does not match orig_val {orig_val:?}, in data set {name}, data `{vals:?}`",
|
||||
);
|
||||
}
|
||||
|
||||
let all_docs: Vec<u32> = (0..vals.len() as u32).collect();
|
||||
buffer.resize(all_docs.len(), 0);
|
||||
reader.get_vals(&all_docs, &mut buffer);
|
||||
assert_eq!(vals, buffer);
|
||||
|
||||
if !vals.is_empty() {
|
||||
let test_rand_idx = rand::thread_rng().gen_range(0..=vals.len() - 1);
|
||||
let expected_positions: Vec<u32> = vals
|
||||
@@ -130,7 +60,7 @@ pub(crate) fn create_and_validate<TColumnCodec: ColumnCodec>(
|
||||
.map(|(pos, _)| pos as u32)
|
||||
.collect();
|
||||
let mut positions = Vec::new();
|
||||
reader.get_row_ids_for_value_range(
|
||||
reader.get_docids_for_value_range(
|
||||
vals[test_rand_idx]..=vals[test_rand_idx],
|
||||
0..vals.len() as u32,
|
||||
&mut positions,
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
use std::fmt::Debug;
|
||||
|
||||
use tantivy_bitpacker::minmax;
|
||||
|
||||
use crate::ColumnValues;
|
||||
|
||||
/// VecColumn provides `Column` over a slice.
|
||||
pub struct VecColumn<'a, T = u64> {
|
||||
pub(crate) values: &'a [T],
|
||||
pub(crate) min_value: T,
|
||||
pub(crate) max_value: T,
|
||||
}
|
||||
|
||||
impl<'a, T: Copy + PartialOrd + Send + Sync + Debug> ColumnValues<T> for VecColumn<'a, T> {
|
||||
fn get_val(&self, position: u32) -> T {
|
||||
self.values[position as usize]
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = T> + '_> {
|
||||
Box::new(self.values.iter().copied())
|
||||
}
|
||||
|
||||
fn min_value(&self) -> T {
|
||||
self.min_value
|
||||
}
|
||||
|
||||
fn max_value(&self) -> T {
|
||||
self.max_value
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.values.len() as u32
|
||||
}
|
||||
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
output.copy_from_slice(&self.values[start as usize..][..output.len()])
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, T: Copy + PartialOrd + Default, V> From<&'a V> for VecColumn<'a, T>
|
||||
where V: AsRef<[T]> + ?Sized
|
||||
{
|
||||
fn from(values: &'a V) -> Self {
|
||||
let values = values.as_ref();
|
||||
let (min_value, max_value) = minmax(values.iter().copied()).unwrap_or_default();
|
||||
Self {
|
||||
values,
|
||||
min_value,
|
||||
max_value,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,44 +1,27 @@
|
||||
use std::fmt;
|
||||
use std::fmt::Debug;
|
||||
use std::net::Ipv6Addr;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::value::NumericalType;
|
||||
use crate::InvalidData;
|
||||
|
||||
/// The column type represents the column type.
|
||||
/// Any changes need to be propagated to `COLUMN_TYPES`.
|
||||
#[derive(Hash, Eq, PartialEq, Debug, Clone, Copy, Ord, PartialOrd, Serialize, Deserialize)]
|
||||
/// The column type represents the column type and can fit on 6-bits.
|
||||
///
|
||||
/// - bits[0..3]: Column category type.
|
||||
/// - bits[3..6]: Numerical type if necessary.
|
||||
#[derive(Hash, Eq, PartialEq, Debug, Clone, Copy, Ord, PartialOrd)]
|
||||
#[repr(u8)]
|
||||
pub enum ColumnType {
|
||||
I64 = 0u8,
|
||||
U64 = 1u8,
|
||||
F64 = 2u8,
|
||||
Bytes = 3u8,
|
||||
Str = 4u8,
|
||||
Bool = 5u8,
|
||||
IpAddr = 6u8,
|
||||
DateTime = 7u8,
|
||||
Bytes = 10u8,
|
||||
Str = 14u8,
|
||||
Bool = 18u8,
|
||||
IpAddr = 22u8,
|
||||
DateTime = 26u8,
|
||||
}
|
||||
|
||||
impl fmt::Display for ColumnType {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
let short_str = match self {
|
||||
ColumnType::I64 => "i64",
|
||||
ColumnType::U64 => "u64",
|
||||
ColumnType::F64 => "f64",
|
||||
ColumnType::Bytes => "bytes",
|
||||
ColumnType::Str => "str",
|
||||
ColumnType::Bool => "bool",
|
||||
ColumnType::IpAddr => "ip",
|
||||
ColumnType::DateTime => "datetime",
|
||||
};
|
||||
write!(f, "{}", short_str)
|
||||
}
|
||||
}
|
||||
|
||||
// The order needs to match _exactly_ the order in the enum
|
||||
#[cfg(test)]
|
||||
const COLUMN_TYPES: [ColumnType; 8] = [
|
||||
ColumnType::I64,
|
||||
ColumnType::U64,
|
||||
@@ -56,7 +39,18 @@ impl ColumnType {
|
||||
}
|
||||
|
||||
pub(crate) fn try_from_code(code: u8) -> Result<ColumnType, InvalidData> {
|
||||
COLUMN_TYPES.get(code as usize).copied().ok_or(InvalidData)
|
||||
use ColumnType::*;
|
||||
match code {
|
||||
0u8 => Ok(I64),
|
||||
1u8 => Ok(U64),
|
||||
2u8 => Ok(F64),
|
||||
10u8 => Ok(Bytes),
|
||||
14u8 => Ok(Str),
|
||||
18u8 => Ok(Bool),
|
||||
22u8 => Ok(IpAddr),
|
||||
26u8 => Ok(Self::DateTime),
|
||||
_ => Err(InvalidData),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -130,7 +124,7 @@ impl HasAssociatedColumnType for bool {
|
||||
}
|
||||
}
|
||||
|
||||
impl HasAssociatedColumnType for common::DateTime {
|
||||
impl HasAssociatedColumnType for crate::DateTime {
|
||||
fn column_type() -> ColumnType {
|
||||
ColumnType::DateTime
|
||||
}
|
||||
@@ -149,20 +143,70 @@ impl HasAssociatedColumnType for Ipv6Addr {
|
||||
}
|
||||
}
|
||||
|
||||
/// Column types are grouped into different categories that
|
||||
/// corresponds to the different types of `JsonValue` types.
|
||||
///
|
||||
/// The columnar writer will apply coercion rules to make sure that
|
||||
/// at most one column exist per `ColumnTypeCategory`.
|
||||
///
|
||||
/// See also [README.md].
|
||||
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
|
||||
#[repr(u8)]
|
||||
pub enum ColumnTypeCategory {
|
||||
Bool,
|
||||
Str,
|
||||
Numerical,
|
||||
DateTime,
|
||||
Bytes,
|
||||
IpAddr,
|
||||
}
|
||||
|
||||
impl From<ColumnType> for ColumnTypeCategory {
|
||||
fn from(column_type: ColumnType) -> Self {
|
||||
match column_type {
|
||||
ColumnType::I64 => ColumnTypeCategory::Numerical,
|
||||
ColumnType::U64 => ColumnTypeCategory::Numerical,
|
||||
ColumnType::F64 => ColumnTypeCategory::Numerical,
|
||||
ColumnType::Bytes => ColumnTypeCategory::Bytes,
|
||||
ColumnType::Str => ColumnTypeCategory::Str,
|
||||
ColumnType::Bool => ColumnTypeCategory::Bool,
|
||||
ColumnType::IpAddr => ColumnTypeCategory::IpAddr,
|
||||
ColumnType::DateTime => ColumnTypeCategory::DateTime,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashSet;
|
||||
|
||||
use super::*;
|
||||
use crate::Cardinality;
|
||||
|
||||
#[test]
|
||||
fn test_column_type_to_code() {
|
||||
for (code, expected_column_type) in super::COLUMN_TYPES.iter().copied().enumerate() {
|
||||
if let Ok(column_type) = ColumnType::try_from_code(code as u8) {
|
||||
assert_eq!(column_type, expected_column_type);
|
||||
let mut column_type_set: HashSet<ColumnType> = HashSet::new();
|
||||
for code in u8::MIN..=u8::MAX {
|
||||
if let Ok(column_type) = ColumnType::try_from_code(code) {
|
||||
assert_eq!(column_type.to_code(), code);
|
||||
assert!(column_type_set.insert(column_type));
|
||||
}
|
||||
}
|
||||
for code in COLUMN_TYPES.len() as u8..=u8::MAX {
|
||||
assert!(ColumnType::try_from_code(code).is_err());
|
||||
assert_eq!(column_type_set.len(), super::COLUMN_TYPES.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_column_category_sort_consistent_with_column_type_sort() {
|
||||
// This is a very important property because we
|
||||
// we need to serialize colunmn in the right order.
|
||||
let mut column_types: Vec<ColumnType> = super::COLUMN_TYPES.iter().copied().collect();
|
||||
column_types.sort_by_key(|col| col.to_code());
|
||||
let column_categories: Vec<ColumnTypeCategory> = column_types
|
||||
.into_iter()
|
||||
.map(ColumnTypeCategory::from)
|
||||
.collect();
|
||||
for (prev, next) in column_categories.iter().zip(column_categories.iter()) {
|
||||
assert!(prev <= next);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ pub const VERSION_FOOTER_NUM_BYTES: usize = MAGIC_BYTES.len() + std::mem::size_o
|
||||
|
||||
/// We end the file by these 4 bytes just to somewhat identify that
|
||||
/// this is indeed a columnar file.
|
||||
const MAGIC_BYTES: [u8; 4] = [2, 113, 119, 66];
|
||||
const MAGIC_BYTES: [u8; 4] = [2, 113, 119, 066];
|
||||
|
||||
pub fn footer() -> [u8; VERSION_FOOTER_NUM_BYTES] {
|
||||
let mut footer_bytes = [0u8; VERSION_FOOTER_NUM_BYTES];
|
||||
@@ -27,8 +27,8 @@ pub enum Version {
|
||||
}
|
||||
|
||||
impl Version {
|
||||
fn to_bytes(self) -> [u8; 4] {
|
||||
(self as u32).to_le_bytes()
|
||||
fn to_bytes(&self) -> [u8; 4] {
|
||||
(*self as u32).to_le_bytes()
|
||||
}
|
||||
|
||||
fn try_from_bytes(bytes: [u8; 4]) -> Result<Version, InvalidData> {
|
||||
|
||||
@@ -1,127 +1,68 @@
|
||||
use std::io::{self, Write};
|
||||
|
||||
use common::{BitSet, CountingWriter, ReadOnlyBitSet};
|
||||
use common::CountingWriter;
|
||||
use sstable::{SSTable, TermOrdinal};
|
||||
|
||||
use super::term_merger::TermMerger;
|
||||
use crate::column::serialize_column_mappable_to_u64;
|
||||
use crate::column_index::SerializableColumnIndex;
|
||||
use crate::iterable::Iterable;
|
||||
use crate::{BytesColumn, MergeRowOrder, ShuffleMergeOrder};
|
||||
use crate::BytesColumn;
|
||||
|
||||
// Serialize [Dictionary, Column, dictionary num bytes U32::LE]
|
||||
// Column: [Column Index, Column Values, column index num bytes U32::LE]
|
||||
pub fn merge_bytes_or_str_column(
|
||||
column_index: SerializableColumnIndex<'_>,
|
||||
bytes_columns: &[Option<BytesColumn>],
|
||||
merge_row_order: &MergeRowOrder,
|
||||
bytes_columns: &[BytesColumn],
|
||||
output: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
// Serialize dict and generate mapping for values
|
||||
let mut output = CountingWriter::wrap(output);
|
||||
// TODO !!! Remove useless terms.
|
||||
let term_ord_mapping = serialize_merged_dict(bytes_columns, merge_row_order, &mut output)?;
|
||||
let term_ord_mapping = serialize_merged_dict(bytes_columns, &mut output)?;
|
||||
let dictionary_num_bytes: u32 = output.written_bytes() as u32;
|
||||
let output = output.finish();
|
||||
let remapped_term_ordinals_values = RemappedTermOrdinalsValues {
|
||||
bytes_columns,
|
||||
term_ord_mapping: &term_ord_mapping,
|
||||
merge_row_order,
|
||||
};
|
||||
serialize_column_mappable_to_u64(column_index, &remapped_term_ordinals_values, output)?;
|
||||
// serialize_bytes_or_str_column(column_index, bytes_columns, &term_ord_mapping, output)?;
|
||||
output.write_all(&dictionary_num_bytes.to_le_bytes())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct RemappedTermOrdinalsValues<'a> {
|
||||
bytes_columns: &'a [Option<BytesColumn>],
|
||||
bytes_columns: &'a [BytesColumn],
|
||||
term_ord_mapping: &'a TermOrdinalMapping,
|
||||
merge_row_order: &'a MergeRowOrder,
|
||||
}
|
||||
|
||||
impl<'a> Iterable for RemappedTermOrdinalsValues<'a> {
|
||||
fn boxed_iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
|
||||
match self.merge_row_order {
|
||||
MergeRowOrder::Stack(_) => self.boxed_iter_stacked(),
|
||||
MergeRowOrder::Shuffled(shuffle_merge_order) => {
|
||||
self.boxed_iter_shuffled(shuffle_merge_order)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> RemappedTermOrdinalsValues<'a> {
|
||||
fn boxed_iter_stacked(&self) -> Box<dyn Iterator<Item = u64> + '_> {
|
||||
let iter = self.bytes_columns.iter().flatten().enumerate().flat_map(
|
||||
move |(seg_ord_with_column, bytes_column)| {
|
||||
let term_ord_after_merge_mapping = self
|
||||
.term_ord_mapping
|
||||
.get_segment(seg_ord_with_column as u32);
|
||||
bytes_column
|
||||
let iter = self
|
||||
.bytes_columns
|
||||
.iter()
|
||||
.enumerate()
|
||||
.flat_map(|(segment_ord, byte_column)| {
|
||||
let segment_ord = self.term_ord_mapping.get_segment(segment_ord);
|
||||
byte_column
|
||||
.ords()
|
||||
.values
|
||||
.iter()
|
||||
.map(move |term_ord| term_ord_after_merge_mapping[term_ord as usize])
|
||||
},
|
||||
);
|
||||
.map(move |term_ord| segment_ord[term_ord as usize])
|
||||
});
|
||||
// TODO see if we can better decompose the mapping / and the stacking
|
||||
Box::new(iter)
|
||||
}
|
||||
|
||||
fn boxed_iter_shuffled<'b>(
|
||||
&'b self,
|
||||
shuffle_merge_order: &'b ShuffleMergeOrder,
|
||||
) -> Box<dyn Iterator<Item = u64> + 'b> {
|
||||
Box::new(
|
||||
shuffle_merge_order
|
||||
.iter_new_to_old_row_addrs()
|
||||
.flat_map(move |old_addr| {
|
||||
let segment_ord = self.term_ord_mapping.get_segment(old_addr.segment_ord);
|
||||
self.bytes_columns[old_addr.segment_ord as usize]
|
||||
.as_ref()
|
||||
.into_iter()
|
||||
.flat_map(move |bytes_column| {
|
||||
bytes_column
|
||||
.term_ords(old_addr.row_id)
|
||||
.map(|old_term_ord: u64| segment_ord[old_term_ord as usize])
|
||||
})
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_term_bitset(column: &BytesColumn, row_bitset: &ReadOnlyBitSet) -> BitSet {
|
||||
let num_terms = column.dictionary().num_terms();
|
||||
let mut term_bitset = BitSet::with_max_value(num_terms as u32);
|
||||
for row_id in row_bitset.iter() {
|
||||
for term_ord in column.term_ord_column.values_for_doc(row_id) {
|
||||
term_bitset.insert(term_ord as u32);
|
||||
}
|
||||
}
|
||||
term_bitset
|
||||
}
|
||||
|
||||
fn is_term_present(bitsets: &[Option<BitSet>], term_merger: &TermMerger) -> bool {
|
||||
for (segment_ord, from_term_ord) in term_merger.matching_segments() {
|
||||
if let Some(bitset) = bitsets[segment_ord].as_ref() {
|
||||
if bitset.contains(from_term_ord as u32) {
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn serialize_merged_dict(
|
||||
bytes_columns: &[Option<BytesColumn>],
|
||||
merge_row_order: &MergeRowOrder,
|
||||
bytes_columns: &[BytesColumn],
|
||||
output: &mut impl Write,
|
||||
) -> io::Result<TermOrdinalMapping> {
|
||||
let mut term_ord_mapping = TermOrdinalMapping::default();
|
||||
|
||||
let mut field_term_streams = Vec::new();
|
||||
for column in bytes_columns.iter().flatten() {
|
||||
for column in bytes_columns {
|
||||
term_ord_mapping.add_segment(column.dictionary.num_terms());
|
||||
let terms = column.dictionary.stream()?;
|
||||
field_term_streams.push(terms);
|
||||
@@ -130,56 +71,21 @@ fn serialize_merged_dict(
|
||||
let mut merged_terms = TermMerger::new(field_term_streams);
|
||||
let mut sstable_builder = sstable::VoidSSTable::writer(output);
|
||||
|
||||
match merge_row_order {
|
||||
MergeRowOrder::Stack(_) => {
|
||||
let mut current_term_ord = 0;
|
||||
while merged_terms.advance() {
|
||||
let term_bytes: &[u8] = merged_terms.key();
|
||||
sstable_builder.insert(term_bytes, &())?;
|
||||
for (segment_ord, from_term_ord) in merged_terms.matching_segments() {
|
||||
term_ord_mapping.register_from_to(segment_ord, from_term_ord, current_term_ord);
|
||||
}
|
||||
current_term_ord += 1;
|
||||
}
|
||||
sstable_builder.finish()?;
|
||||
}
|
||||
MergeRowOrder::Shuffled(shuffle_merge_order) => {
|
||||
assert_eq!(shuffle_merge_order.alive_bitsets.len(), bytes_columns.len());
|
||||
let mut term_bitsets: Vec<Option<BitSet>> = Vec::with_capacity(bytes_columns.len());
|
||||
for (alive_bitset_opt, bytes_column_opt) in shuffle_merge_order
|
||||
.alive_bitsets
|
||||
.iter()
|
||||
.zip(bytes_columns.iter())
|
||||
{
|
||||
match (alive_bitset_opt, bytes_column_opt) {
|
||||
(Some(alive_bitset), Some(bytes_column)) => {
|
||||
let term_bitset = compute_term_bitset(bytes_column, alive_bitset);
|
||||
term_bitsets.push(Some(term_bitset));
|
||||
}
|
||||
_ => {
|
||||
term_bitsets.push(None);
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut current_term_ord = 0;
|
||||
while merged_terms.advance() {
|
||||
let term_bytes: &[u8] = merged_terms.key();
|
||||
if !is_term_present(&term_bitsets[..], &merged_terms) {
|
||||
continue;
|
||||
}
|
||||
sstable_builder.insert(term_bytes, &())?;
|
||||
for (segment_ord, from_term_ord) in merged_terms.matching_segments() {
|
||||
term_ord_mapping.register_from_to(segment_ord, from_term_ord, current_term_ord);
|
||||
}
|
||||
current_term_ord += 1;
|
||||
}
|
||||
sstable_builder.finish()?;
|
||||
let mut current_term_ord = 0;
|
||||
while merged_terms.advance() {
|
||||
let term_bytes: &[u8] = merged_terms.key();
|
||||
|
||||
sstable_builder.insert(term_bytes, &())?;
|
||||
for (segment_ord, from_term_ord) in merged_terms.matching_segments() {
|
||||
term_ord_mapping.register_from_to(segment_ord, from_term_ord, current_term_ord);
|
||||
}
|
||||
current_term_ord += 1;
|
||||
}
|
||||
sstable_builder.finish()?;
|
||||
Ok(term_ord_mapping)
|
||||
}
|
||||
|
||||
#[derive(Default, Debug)]
|
||||
#[derive(Default)]
|
||||
struct TermOrdinalMapping {
|
||||
per_segment_new_term_ordinals: Vec<Vec<TermOrdinal>>,
|
||||
}
|
||||
@@ -187,14 +93,14 @@ struct TermOrdinalMapping {
|
||||
impl TermOrdinalMapping {
|
||||
fn add_segment(&mut self, max_term_ord: usize) {
|
||||
self.per_segment_new_term_ordinals
|
||||
.push(vec![TermOrdinal::default(); max_term_ord]);
|
||||
.push(vec![TermOrdinal::default(); max_term_ord as usize]);
|
||||
}
|
||||
|
||||
fn register_from_to(&mut self, segment_ord: usize, from_ord: TermOrdinal, to_ord: TermOrdinal) {
|
||||
self.per_segment_new_term_ordinals[segment_ord][from_ord as usize] = to_ord;
|
||||
}
|
||||
|
||||
fn get_segment(&self, segment_ord: u32) -> &[TermOrdinal] {
|
||||
&(self.per_segment_new_term_ordinals[segment_ord as usize])[..]
|
||||
fn get_segment(&self, segment_ord: usize) -> &[TermOrdinal] {
|
||||
&(self.per_segment_new_term_ordinals[segment_ord])[..]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
use std::ops::Range;
|
||||
|
||||
use common::{BitSet, OwnedBytes, ReadOnlyBitSet};
|
||||
|
||||
use crate::{ColumnarReader, RowAddr, RowId};
|
||||
use crate::{column, ColumnarReader, RowId};
|
||||
|
||||
pub struct StackMergeOrder {
|
||||
// This does not start at 0. The first row is the number of
|
||||
@@ -44,75 +42,19 @@ pub enum MergeRowOrder {
|
||||
/// rows [r0..n_row_0) contains the row of columnar_readers[0], in ordder
|
||||
/// rows [n_row_0..n_row_0 + n_row_1 contains the row of columnar_readers[1], in order.
|
||||
/// ..
|
||||
/// No documents is deleted.
|
||||
Stack(StackMergeOrder),
|
||||
/// Some more complex mapping, that may interleaves rows from the different readers and
|
||||
/// drop rows, or do both.
|
||||
Shuffled(ShuffleMergeOrder),
|
||||
}
|
||||
|
||||
impl From<StackMergeOrder> for MergeRowOrder {
|
||||
fn from(stack_merge_order: StackMergeOrder) -> MergeRowOrder {
|
||||
MergeRowOrder::Stack(stack_merge_order)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ShuffleMergeOrder> for MergeRowOrder {
|
||||
fn from(shuffle_merge_order: ShuffleMergeOrder) -> MergeRowOrder {
|
||||
MergeRowOrder::Shuffled(shuffle_merge_order)
|
||||
}
|
||||
/// Some more complex mapping, that can interleaves rows from the different readers and
|
||||
/// possibly drop rows.
|
||||
Complex(()),
|
||||
}
|
||||
|
||||
impl MergeRowOrder {
|
||||
pub fn num_rows(&self) -> RowId {
|
||||
match self {
|
||||
MergeRowOrder::Stack(stack_row_order) => stack_row_order.num_rows(),
|
||||
MergeRowOrder::Shuffled(complex_mapping) => complex_mapping.num_rows(),
|
||||
MergeRowOrder::Complex(_) => {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ShuffleMergeOrder {
|
||||
pub new_row_id_to_old_row_id: Vec<RowAddr>,
|
||||
pub alive_bitsets: Vec<Option<ReadOnlyBitSet>>,
|
||||
}
|
||||
|
||||
impl ShuffleMergeOrder {
|
||||
pub fn for_test(
|
||||
segment_num_rows: &[RowId],
|
||||
new_row_id_to_old_row_id: Vec<RowAddr>,
|
||||
) -> ShuffleMergeOrder {
|
||||
let mut alive_bitsets: Vec<BitSet> = segment_num_rows
|
||||
.iter()
|
||||
.map(|&num_rows| BitSet::with_max_value(num_rows))
|
||||
.collect();
|
||||
for &RowAddr {
|
||||
segment_ord,
|
||||
row_id,
|
||||
} in &new_row_id_to_old_row_id
|
||||
{
|
||||
alive_bitsets[segment_ord as usize].insert(row_id);
|
||||
}
|
||||
let alive_bitsets: Vec<Option<ReadOnlyBitSet>> = alive_bitsets
|
||||
.into_iter()
|
||||
.map(|alive_bitset| {
|
||||
let mut buffer = Vec::new();
|
||||
alive_bitset.serialize(&mut buffer).unwrap();
|
||||
let data = OwnedBytes::new(buffer);
|
||||
Some(ReadOnlyBitSet::open(data))
|
||||
})
|
||||
.collect();
|
||||
ShuffleMergeOrder {
|
||||
new_row_id_to_old_row_id,
|
||||
alive_bitsets,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn num_rows(&self) -> RowId {
|
||||
self.new_row_id_to_old_row_id.len() as RowId
|
||||
}
|
||||
|
||||
pub fn iter_new_to_old_row_addrs(&self) -> impl Iterator<Item = RowAddr> + '_ {
|
||||
self.new_row_id_to_old_row_id.iter().copied()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,100 +2,42 @@ mod merge_dict_column;
|
||||
mod merge_mapping;
|
||||
mod term_merger;
|
||||
|
||||
// mod sorted_doc_id_column;
|
||||
|
||||
use std::collections::{BTreeMap, HashMap, HashSet};
|
||||
use std::io;
|
||||
use std::net::Ipv6Addr;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use merge_mapping::{MergeRowOrder, ShuffleMergeOrder, StackMergeOrder};
|
||||
pub use merge_mapping::{MergeRowOrder, StackMergeOrder};
|
||||
|
||||
use super::writer::ColumnarSerializer;
|
||||
use crate::column::{serialize_column_mappable_to_u128, serialize_column_mappable_to_u64};
|
||||
use crate::column_values::MergedColumnValues;
|
||||
use crate::column_index::stack_column_index;
|
||||
use crate::columnar::column_type::ColumnTypeCategory;
|
||||
use crate::columnar::merge::merge_dict_column::merge_bytes_or_str_column;
|
||||
use crate::columnar::writer::CompatibleNumericalTypes;
|
||||
use crate::columnar::ColumnarReader;
|
||||
use crate::dynamic_column::DynamicColumn;
|
||||
use crate::{
|
||||
BytesColumn, Column, ColumnIndex, ColumnType, ColumnValues, NumericalType, NumericalValue,
|
||||
BytesColumn, Column, ColumnIndex, ColumnType, ColumnValues, MonotonicallyMappableToU128,
|
||||
NumericalType, NumericalValue,
|
||||
};
|
||||
|
||||
/// Column types are grouped into different categories.
|
||||
/// After merge, all columns belonging to the same category are coerced to
|
||||
/// the same column type.
|
||||
///
|
||||
/// In practise, today, only Numerical colummns are coerced into one type today.
|
||||
///
|
||||
/// See also [README.md].
|
||||
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
|
||||
pub(crate) enum ColumnTypeCategory {
|
||||
Bool,
|
||||
Str,
|
||||
Numerical,
|
||||
DateTime,
|
||||
Bytes,
|
||||
IpAddr,
|
||||
}
|
||||
|
||||
impl From<ColumnType> for ColumnTypeCategory {
|
||||
fn from(column_type: ColumnType) -> Self {
|
||||
match column_type {
|
||||
ColumnType::I64 => ColumnTypeCategory::Numerical,
|
||||
ColumnType::U64 => ColumnTypeCategory::Numerical,
|
||||
ColumnType::F64 => ColumnTypeCategory::Numerical,
|
||||
ColumnType::Bytes => ColumnTypeCategory::Bytes,
|
||||
ColumnType::Str => ColumnTypeCategory::Str,
|
||||
ColumnType::Bool => ColumnTypeCategory::Bool,
|
||||
ColumnType::IpAddr => ColumnTypeCategory::IpAddr,
|
||||
ColumnType::DateTime => ColumnTypeCategory::DateTime,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Merge several columnar table together.
|
||||
///
|
||||
/// If several columns with the same name are conflicting with the numerical types in the
|
||||
/// input columnars, the first type compatible out of i64, u64, f64 in that order will be used.
|
||||
///
|
||||
/// `require_columns` makes it possible to ensure that some columns will be present in the
|
||||
/// resulting columnar. When a required column is a numerical column type, one of two things can
|
||||
/// happen:
|
||||
/// - If the required column type is compatible with all of the input columnar, the resulsting
|
||||
/// merged
|
||||
/// columnar will simply coerce the input column and use the required column type.
|
||||
/// - If the required column type is incompatible with one of the input columnar, the merged
|
||||
/// will fail with an InvalidData error.
|
||||
///
|
||||
/// `merge_row_order` makes it possible to remove or reorder row in the resulting
|
||||
/// `Columnar` table.
|
||||
///
|
||||
/// Reminder: a string and a numerical column may bare the same column name. This is not
|
||||
/// considered a conflict.
|
||||
pub fn merge_columnar(
|
||||
columnar_readers: &[&ColumnarReader],
|
||||
required_columns: &[(String, ColumnType)],
|
||||
merge_row_order: MergeRowOrder,
|
||||
mapping: MergeRowOrder,
|
||||
output: &mut impl io::Write,
|
||||
) -> io::Result<()> {
|
||||
let mut serializer = ColumnarSerializer::new(output);
|
||||
let num_rows_per_columnar = columnar_readers
|
||||
.iter()
|
||||
.map(|reader| reader.num_rows())
|
||||
.collect::<Vec<u32>>();
|
||||
|
||||
let columns_to_merge = group_columns_for_merge(columnar_readers, required_columns)?;
|
||||
let columns_to_merge = group_columns_for_merge(columnar_readers)?;
|
||||
for ((column_name, column_type), columns) in columns_to_merge {
|
||||
let mut column_serializer =
|
||||
serializer.serialize_column(column_name.as_bytes(), column_type);
|
||||
merge_column(
|
||||
column_type,
|
||||
&num_rows_per_columnar,
|
||||
columns,
|
||||
&merge_row_order,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
merge_column(column_type, columns, &mapping, &mut column_serializer)?;
|
||||
}
|
||||
serializer.finalize(merge_row_order.num_rows())?;
|
||||
serializer.finalize(mapping.num_rows())?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -111,9 +53,8 @@ fn dynamic_column_to_u64_monotonic(dynamic_column: DynamicColumn) -> Option<Colu
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_column(
|
||||
pub fn merge_column(
|
||||
column_type: ColumnType,
|
||||
num_docs_per_column: &[u32],
|
||||
columns: Vec<Option<DynamicColumn>>,
|
||||
merge_row_order: &MergeRowOrder,
|
||||
wrt: &mut impl io::Write,
|
||||
@@ -124,265 +65,160 @@ fn merge_column(
|
||||
| ColumnType::F64
|
||||
| ColumnType::DateTime
|
||||
| ColumnType::Bool => {
|
||||
let mut column_indexes: Vec<ColumnIndex> = Vec::with_capacity(columns.len());
|
||||
let mut column_values: Vec<Option<Arc<dyn ColumnValues>>> =
|
||||
Vec::with_capacity(columns.len());
|
||||
for (i, dynamic_column_opt) in columns.into_iter().enumerate() {
|
||||
if let Some(Column { index: idx, values }) =
|
||||
let mut column_indexes: Vec<Option<ColumnIndex>> = Vec::with_capacity(columns.len());
|
||||
let mut column_values: Vec<Arc<dyn ColumnValues>> = Vec::with_capacity(columns.len());
|
||||
for dynamic_column_opt in columns {
|
||||
if let Some(Column { idx, values }) =
|
||||
dynamic_column_opt.and_then(dynamic_column_to_u64_monotonic)
|
||||
{
|
||||
column_indexes.push(idx);
|
||||
column_values.push(Some(values));
|
||||
column_indexes.push(Some(idx));
|
||||
column_values.push(values);
|
||||
} else {
|
||||
column_indexes.push(ColumnIndex::Empty {
|
||||
num_docs: num_docs_per_column[i],
|
||||
});
|
||||
column_values.push(None);
|
||||
column_indexes.push(None);
|
||||
}
|
||||
}
|
||||
let merged_column_index =
|
||||
crate::column_index::merge_column_index(&column_indexes[..], merge_row_order);
|
||||
let merge_column_values = MergedColumnValues {
|
||||
column_indexes: &column_indexes[..],
|
||||
column_values: &column_values[..],
|
||||
merge_row_order,
|
||||
};
|
||||
serialize_column_mappable_to_u64(merged_column_index, &merge_column_values, wrt)?;
|
||||
crate::column_index::stack_column_index(&column_indexes[..], merge_row_order);
|
||||
serialize_column_mappable_to_u64(merged_column_index, &&column_values[..], wrt)?;
|
||||
}
|
||||
ColumnType::IpAddr => {
|
||||
let mut column_indexes: Vec<ColumnIndex> = Vec::with_capacity(columns.len());
|
||||
let mut column_values: Vec<Option<Arc<dyn ColumnValues<Ipv6Addr>>>> =
|
||||
let mut column_indexes: Vec<Option<ColumnIndex>> = Vec::with_capacity(columns.len());
|
||||
let mut column_values: Vec<Arc<dyn ColumnValues<Ipv6Addr>>> =
|
||||
Vec::with_capacity(columns.len());
|
||||
for (i, dynamic_column_opt) in columns.into_iter().enumerate() {
|
||||
if let Some(DynamicColumn::IpAddr(Column { index: idx, values })) =
|
||||
dynamic_column_opt
|
||||
{
|
||||
column_indexes.push(idx);
|
||||
column_values.push(Some(values));
|
||||
let mut num_values = 0;
|
||||
for dynamic_column_opt in columns {
|
||||
if let Some(DynamicColumn::IpAddr(Column { idx, values })) = dynamic_column_opt {
|
||||
num_values += values.num_vals();
|
||||
column_indexes.push(Some(idx));
|
||||
column_values.push(values);
|
||||
} else {
|
||||
column_indexes.push(ColumnIndex::Empty {
|
||||
num_docs: num_docs_per_column[i],
|
||||
});
|
||||
column_values.push(None);
|
||||
column_indexes.push(None);
|
||||
}
|
||||
}
|
||||
|
||||
let merged_column_index =
|
||||
crate::column_index::merge_column_index(&column_indexes[..], merge_row_order);
|
||||
let merge_column_values = MergedColumnValues {
|
||||
column_indexes: &column_indexes[..],
|
||||
column_values: &column_values,
|
||||
merge_row_order,
|
||||
};
|
||||
|
||||
serialize_column_mappable_to_u128(merged_column_index, &merge_column_values, wrt)?;
|
||||
crate::column_index::stack_column_index(&column_indexes[..], merge_row_order);
|
||||
serialize_column_mappable_to_u128(
|
||||
merged_column_index,
|
||||
&&column_values[..],
|
||||
num_values,
|
||||
wrt,
|
||||
)?;
|
||||
}
|
||||
ColumnType::Bytes | ColumnType::Str => {
|
||||
let mut column_indexes: Vec<ColumnIndex> = Vec::with_capacity(columns.len());
|
||||
let mut bytes_columns: Vec<Option<BytesColumn>> = Vec::with_capacity(columns.len());
|
||||
for (i, dynamic_column_opt) in columns.into_iter().enumerate() {
|
||||
let mut column_indexes: Vec<Option<ColumnIndex>> = Vec::with_capacity(columns.len());
|
||||
let mut bytes_columns: Vec<BytesColumn> = Vec::with_capacity(columns.len());
|
||||
for dynamic_column_opt in columns {
|
||||
match dynamic_column_opt {
|
||||
Some(DynamicColumn::Str(str_column)) => {
|
||||
column_indexes.push(str_column.term_ord_column.index.clone());
|
||||
bytes_columns.push(Some(str_column.into()));
|
||||
column_indexes.push(Some(str_column.term_ord_column.idx.clone()));
|
||||
bytes_columns.push(str_column.into());
|
||||
}
|
||||
Some(DynamicColumn::Bytes(bytes_column)) => {
|
||||
column_indexes.push(bytes_column.term_ord_column.index.clone());
|
||||
bytes_columns.push(Some(bytes_column));
|
||||
}
|
||||
_ => {
|
||||
column_indexes.push(ColumnIndex::Empty {
|
||||
num_docs: num_docs_per_column[i],
|
||||
});
|
||||
bytes_columns.push(None);
|
||||
column_indexes.push(Some(bytes_column.term_ord_column.idx.clone()));
|
||||
bytes_columns.push(bytes_column);
|
||||
}
|
||||
_ => column_indexes.push(None),
|
||||
}
|
||||
}
|
||||
let merged_column_index =
|
||||
crate::column_index::merge_column_index(&column_indexes[..], merge_row_order);
|
||||
merge_bytes_or_str_column(merged_column_index, &bytes_columns, merge_row_order, wrt)?;
|
||||
crate::column_index::stack_column_index(&column_indexes[..], merge_row_order);
|
||||
merge_bytes_or_str_column(merged_column_index, &bytes_columns, wrt)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct GroupedColumns {
|
||||
required_column_type: Option<ColumnType>,
|
||||
columns: Vec<Option<DynamicColumn>>,
|
||||
column_category: ColumnTypeCategory,
|
||||
}
|
||||
|
||||
impl GroupedColumns {
|
||||
fn for_category(column_category: ColumnTypeCategory, num_columnars: usize) -> Self {
|
||||
GroupedColumns {
|
||||
required_column_type: None,
|
||||
columns: vec![None; num_columnars],
|
||||
column_category,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the dynamic column for a given columnar.
|
||||
fn set_column(&mut self, columnar_id: usize, column: DynamicColumn) {
|
||||
self.columns[columnar_id] = Some(column);
|
||||
}
|
||||
|
||||
/// Force the existence of a column, as well as its type.
|
||||
fn require_type(&mut self, required_type: ColumnType) -> io::Result<()> {
|
||||
if let Some(existing_required_type) = self.required_column_type {
|
||||
if existing_required_type == required_type {
|
||||
// This was just a duplicate in the `required_columns`.
|
||||
// Nothing to do.
|
||||
return Ok(());
|
||||
} else {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"Required column conflicts with another required column of the same type \
|
||||
category.",
|
||||
));
|
||||
}
|
||||
}
|
||||
self.required_column_type = Some(required_type);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the column type after merge.
|
||||
///
|
||||
/// This method does not check if the column types can actually be coerced to
|
||||
/// this type.
|
||||
fn column_type_after_merge(&self) -> ColumnType {
|
||||
if let Some(required_type) = self.required_column_type {
|
||||
return required_type;
|
||||
}
|
||||
let column_type: HashSet<ColumnType> = self
|
||||
.columns
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(|column| column.column_type())
|
||||
.collect();
|
||||
if column_type.len() == 1 {
|
||||
return column_type.into_iter().next().unwrap();
|
||||
}
|
||||
// At the moment, only the numerical categorical column type has more than one possible
|
||||
// column type.
|
||||
assert_eq!(self.column_category, ColumnTypeCategory::Numerical);
|
||||
merged_numerical_columns_type(self.columns.iter().flatten()).into()
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the type of the merged numerical column.
|
||||
///
|
||||
/// This function picks the first numerical type out of i64, u64, f64 (order matters
|
||||
/// here), that is compatible with all the `columns`.
|
||||
///
|
||||
/// # Panics
|
||||
/// Panics if one of the column is not numerical.
|
||||
fn merged_numerical_columns_type<'a>(
|
||||
columns: impl Iterator<Item = &'a DynamicColumn>,
|
||||
) -> NumericalType {
|
||||
let mut compatible_numerical_types = CompatibleNumericalTypes::default();
|
||||
for column in columns {
|
||||
let (min_value, max_value) =
|
||||
min_max_if_numerical(column).expect("All columns re required to be numerical");
|
||||
compatible_numerical_types.accept_value(min_value);
|
||||
compatible_numerical_types.accept_value(max_value);
|
||||
}
|
||||
compatible_numerical_types.to_numerical_type()
|
||||
}
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn group_columns_for_merge(
|
||||
columnar_readers: &[&ColumnarReader],
|
||||
required_columns: &[(String, ColumnType)],
|
||||
) -> io::Result<BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>>> {
|
||||
// Each column name may have multiple types of column associated.
|
||||
// For merging we are interested in the same column type category since they can be merged.
|
||||
let mut columns_grouped: HashMap<(String, ColumnTypeCategory), GroupedColumns> = HashMap::new();
|
||||
let mut columns_grouped: HashMap<(String, ColumnTypeCategory), Vec<Option<DynamicColumn>>> =
|
||||
HashMap::new();
|
||||
|
||||
for &(ref column_name, column_type) in required_columns {
|
||||
columns_grouped
|
||||
.entry((column_name.clone(), column_type.into()))
|
||||
.or_insert_with(|| {
|
||||
GroupedColumns::for_category(column_type.into(), columnar_readers.len())
|
||||
})
|
||||
.require_type(column_type)?;
|
||||
}
|
||||
let num_columnars = columnar_readers.len();
|
||||
|
||||
for (columnar_id, columnar_reader) in columnar_readers.iter().enumerate() {
|
||||
let column_name_and_handle = columnar_reader.list_columns()?;
|
||||
for (column_name, handle) in column_name_and_handle {
|
||||
let column_category: ColumnTypeCategory = handle.column_type().into();
|
||||
let column_type_category: ColumnTypeCategory = handle.column_type().into();
|
||||
let columns = columns_grouped
|
||||
.entry((column_name, column_type_category))
|
||||
.or_insert_with(|| vec![None; num_columnars]);
|
||||
let column = handle.open()?;
|
||||
columns_grouped
|
||||
.entry((column_name, column_category))
|
||||
.or_insert_with(|| {
|
||||
GroupedColumns::for_category(column_category, columnar_readers.len())
|
||||
})
|
||||
.set_column(columnar_id, column);
|
||||
columns[columnar_id] = Some(column);
|
||||
}
|
||||
}
|
||||
|
||||
let mut merge_columns: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
Default::default();
|
||||
BTreeMap::default();
|
||||
|
||||
for ((column_name, _), mut grouped_columns) in columns_grouped {
|
||||
let column_type = grouped_columns.column_type_after_merge();
|
||||
coerce_columns(column_type, &mut grouped_columns.columns)?;
|
||||
merge_columns.insert((column_name, column_type), grouped_columns.columns);
|
||||
for ((column_name, col_category), mut columns) in columns_grouped {
|
||||
if col_category == ColumnTypeCategory::Numerical {
|
||||
coerce_numerical_columns_to_same_type(&mut columns);
|
||||
}
|
||||
let column_type = columns
|
||||
.iter()
|
||||
.flatten()
|
||||
.map(|col| col.column_type())
|
||||
.next()
|
||||
.unwrap();
|
||||
merge_columns.insert((column_name, column_type), columns);
|
||||
}
|
||||
|
||||
Ok(merge_columns)
|
||||
}
|
||||
|
||||
fn coerce_columns(
|
||||
column_type: ColumnType,
|
||||
columns: &mut [Option<DynamicColumn>],
|
||||
) -> io::Result<()> {
|
||||
/// Coerce a set of numerical columns to the same type.
|
||||
///
|
||||
/// If all columns are already from the same type, keep this type
|
||||
/// (even if they could all be coerced to i64).
|
||||
fn coerce_numerical_columns_to_same_type(columns: &mut [Option<DynamicColumn>]) {
|
||||
let mut column_types: HashSet<NumericalType> = HashSet::default();
|
||||
let mut compatible_numerical_types = CompatibleNumericalTypes::default();
|
||||
for column in columns.iter().flatten() {
|
||||
let min_value: NumericalValue;
|
||||
let max_value: NumericalValue;
|
||||
match column {
|
||||
DynamicColumn::I64(column) => {
|
||||
min_value = column.min_value().into();
|
||||
max_value = column.max_value().into();
|
||||
}
|
||||
DynamicColumn::U64(column) => {
|
||||
min_value = column.min_value().into();
|
||||
max_value = column.min_value().into();
|
||||
}
|
||||
DynamicColumn::F64(column) => {
|
||||
min_value = column.min_value().into();
|
||||
max_value = column.min_value().into();
|
||||
}
|
||||
DynamicColumn::Bool(_)
|
||||
| DynamicColumn::IpAddr(_)
|
||||
| DynamicColumn::DateTime(_)
|
||||
| DynamicColumn::Bytes(_)
|
||||
| DynamicColumn::Str(_) => {
|
||||
panic!("We expected only numerical columns.");
|
||||
}
|
||||
}
|
||||
column_types.insert(column.column_type().numerical_type().unwrap());
|
||||
compatible_numerical_types.accept_value(min_value);
|
||||
compatible_numerical_types.accept_value(max_value);
|
||||
}
|
||||
if column_types.len() <= 1 {
|
||||
// No need to do anything. The columns are already all from the same type.
|
||||
// This is necessary to let use force a given type.
|
||||
|
||||
// TODO This works in a world where we do not allow a change of schema,
|
||||
// but in the future, we will have to pass some kind of schema to enforce
|
||||
// the logic.
|
||||
return;
|
||||
}
|
||||
let coerce_type = compatible_numerical_types.to_numerical_type();
|
||||
for column_opt in columns.iter_mut() {
|
||||
if let Some(column) = column_opt.take() {
|
||||
*column_opt = Some(coerce_column(column_type, column)?);
|
||||
*column_opt = column.coerce_numerical(coerce_type);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn coerce_column(column_type: ColumnType, column: DynamicColumn) -> io::Result<DynamicColumn> {
|
||||
if let Some(numerical_type) = column_type.numerical_type() {
|
||||
column
|
||||
.coerce_numerical(numerical_type)
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, ""))
|
||||
} else {
|
||||
if column.column_type() != column_type {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"Cannot coerce column of type `{:?}` to `{column_type:?}`",
|
||||
column.column_type()
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(column)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the (min, max) of a column provided it is numerical (i64, u64. f64).
|
||||
///
|
||||
/// The min and the max are simply the numerical value as defined by `ColumnValue::min_value()`,
|
||||
/// and `ColumnValue::max_value()`.
|
||||
///
|
||||
/// It is important to note that these values are only guaranteed to be lower/upper bound
|
||||
/// (as opposed to min/max value).
|
||||
/// If a column is empty, the min and max values are currently set to 0.
|
||||
fn min_max_if_numerical(column: &DynamicColumn) -> Option<(NumericalValue, NumericalValue)> {
|
||||
match column {
|
||||
DynamicColumn::I64(column) => Some((column.min_value().into(), column.max_value().into())),
|
||||
DynamicColumn::U64(column) => Some((column.min_value().into(), column.min_value().into())),
|
||||
DynamicColumn::F64(column) => Some((column.min_value().into(), column.min_value().into())),
|
||||
DynamicColumn::Bool(_)
|
||||
| DynamicColumn::IpAddr(_)
|
||||
| DynamicColumn::DateTime(_)
|
||||
| DynamicColumn::Bytes(_)
|
||||
| DynamicColumn::Str(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
107
columnar/src/columnar/merge/sorted_doc_id_column.rs
Normal file
107
columnar/src/columnar/merge/sorted_doc_id_column.rs
Normal file
@@ -0,0 +1,107 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use fastfield_codecs::Column;
|
||||
use itertools::Itertools;
|
||||
|
||||
use crate::indexer::doc_id_mapping::SegmentDocIdMapping;
|
||||
use crate::SegmentReader;
|
||||
|
||||
pub(crate) struct RemappedDocIdColumn<'a> {
|
||||
doc_id_mapping: &'a SegmentDocIdMapping,
|
||||
fast_field_readers: Vec<Arc<dyn Column<u64>>>,
|
||||
min_value: u64,
|
||||
max_value: u64,
|
||||
num_vals: u32,
|
||||
}
|
||||
|
||||
fn compute_min_max_val(
|
||||
u64_reader: &dyn Column<u64>,
|
||||
segment_reader: &SegmentReader,
|
||||
) -> Option<(u64, u64)> {
|
||||
if segment_reader.max_doc() == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
if segment_reader.alive_bitset().is_none() {
|
||||
// no deleted documents,
|
||||
// we can use the previous min_val, max_val.
|
||||
return Some((u64_reader.min_value(), u64_reader.max_value()));
|
||||
}
|
||||
// some deleted documents,
|
||||
// we need to recompute the max / min
|
||||
segment_reader
|
||||
.doc_ids_alive()
|
||||
.map(|doc_id| u64_reader.get_val(doc_id))
|
||||
.minmax()
|
||||
.into_option()
|
||||
}
|
||||
|
||||
impl<'a> RemappedDocIdColumn<'a> {
|
||||
pub(crate) fn new(
|
||||
readers: &'a [SegmentReader],
|
||||
doc_id_mapping: &'a SegmentDocIdMapping,
|
||||
field: &str,
|
||||
) -> Self {
|
||||
let (min_value, max_value) = readers
|
||||
.iter()
|
||||
.filter_map(|reader| {
|
||||
let u64_reader: Arc<dyn Column<u64>> =
|
||||
reader.fast_fields().typed_fast_field_reader(field).expect(
|
||||
"Failed to find a reader for single fast field. This is a tantivy bug and \
|
||||
it should never happen.",
|
||||
);
|
||||
compute_min_max_val(&*u64_reader, reader)
|
||||
})
|
||||
.reduce(|a, b| (a.0.min(b.0), a.1.max(b.1)))
|
||||
.expect("Unexpected error, empty readers in IndexMerger");
|
||||
|
||||
let fast_field_readers = readers
|
||||
.iter()
|
||||
.map(|reader| {
|
||||
let u64_reader: Arc<dyn Column<u64>> =
|
||||
reader.fast_fields().typed_fast_field_reader(field).expect(
|
||||
"Failed to find a reader for single fast field. This is a tantivy bug and \
|
||||
it should never happen.",
|
||||
);
|
||||
u64_reader
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
RemappedDocIdColumn {
|
||||
doc_id_mapping,
|
||||
fast_field_readers,
|
||||
min_value,
|
||||
max_value,
|
||||
num_vals: doc_id_mapping.len() as u32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Column for RemappedDocIdColumn<'a> {
|
||||
fn get_val(&self, _doc: u32) -> u64 {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
|
||||
Box::new(
|
||||
self.doc_id_mapping
|
||||
.iter_old_doc_addrs()
|
||||
.map(|old_doc_addr| {
|
||||
let fast_field_reader =
|
||||
&self.fast_field_readers[old_doc_addr.segment_ord as usize];
|
||||
fast_field_reader.get_val(old_doc_addr.doc_id)
|
||||
}),
|
||||
)
|
||||
}
|
||||
fn min_value(&self) -> u64 {
|
||||
self.min_value
|
||||
}
|
||||
|
||||
fn max_value(&self) -> u64 {
|
||||
self.max_value
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.num_vals
|
||||
}
|
||||
}
|
||||
169
columnar/src/columnar/merge/sorted_doc_id_multivalue_column.rs
Normal file
169
columnar/src/columnar/merge/sorted_doc_id_multivalue_column.rs
Normal file
@@ -0,0 +1,169 @@
|
||||
use std::cmp;
|
||||
|
||||
use fastfield_codecs::Column;
|
||||
|
||||
use super::flat_map_with_buffer::FlatMapWithBufferIter;
|
||||
use crate::fastfield::{MultiValueIndex, MultiValuedFastFieldReader};
|
||||
use crate::indexer::doc_id_mapping::SegmentDocIdMapping;
|
||||
use crate::{DocAddress, SegmentReader};
|
||||
|
||||
pub(crate) struct RemappedDocIdMultiValueColumn<'a> {
|
||||
doc_id_mapping: &'a SegmentDocIdMapping,
|
||||
fast_field_readers: Vec<MultiValuedFastFieldReader<u64>>,
|
||||
min_value: u64,
|
||||
max_value: u64,
|
||||
num_vals: u32,
|
||||
}
|
||||
|
||||
impl<'a> RemappedDocIdMultiValueColumn<'a> {
|
||||
pub(crate) fn new(
|
||||
readers: &'a [SegmentReader],
|
||||
doc_id_mapping: &'a SegmentDocIdMapping,
|
||||
field: &str,
|
||||
) -> Self {
|
||||
// Our values are bitpacked and we need to know what should be
|
||||
// our bitwidth and our minimum value before serializing any values.
|
||||
//
|
||||
// Computing those is non-trivial if some documents are deleted.
|
||||
// We go through a complete first pass to compute the minimum and the
|
||||
// maximum value and initialize our Serializer.
|
||||
let mut num_vals = 0;
|
||||
let mut min_value = u64::MAX;
|
||||
let mut max_value = u64::MIN;
|
||||
let mut vals = Vec::new();
|
||||
let mut fast_field_readers = Vec::with_capacity(readers.len());
|
||||
for reader in readers {
|
||||
let ff_reader: MultiValuedFastFieldReader<u64> = reader
|
||||
.fast_fields()
|
||||
.typed_fast_field_multi_reader::<u64>(field)
|
||||
.expect(
|
||||
"Failed to find multivalued fast field reader. This is a bug in tantivy. \
|
||||
Please report.",
|
||||
);
|
||||
for doc in reader.doc_ids_alive() {
|
||||
ff_reader.get_vals(doc, &mut vals);
|
||||
for &val in &vals {
|
||||
min_value = cmp::min(val, min_value);
|
||||
max_value = cmp::max(val, max_value);
|
||||
}
|
||||
num_vals += vals.len();
|
||||
}
|
||||
fast_field_readers.push(ff_reader);
|
||||
// TODO optimize when no deletes
|
||||
}
|
||||
if min_value > max_value {
|
||||
min_value = 0;
|
||||
max_value = 0;
|
||||
}
|
||||
RemappedDocIdMultiValueColumn {
|
||||
doc_id_mapping,
|
||||
fast_field_readers,
|
||||
min_value,
|
||||
max_value,
|
||||
num_vals: num_vals as u32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Column for RemappedDocIdMultiValueColumn<'a> {
|
||||
fn get_val(&self, _pos: u32) -> u64 {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
|
||||
Box::new(
|
||||
self.doc_id_mapping
|
||||
.iter_old_doc_addrs()
|
||||
.flat_map_with_buffer(|old_doc_addr: DocAddress, buffer| {
|
||||
let ff_reader = &self.fast_field_readers[old_doc_addr.segment_ord as usize];
|
||||
ff_reader.get_vals(old_doc_addr.doc_id, buffer);
|
||||
}),
|
||||
)
|
||||
}
|
||||
fn min_value(&self) -> u64 {
|
||||
self.min_value
|
||||
}
|
||||
|
||||
fn max_value(&self) -> u64 {
|
||||
self.max_value
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.num_vals
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct RemappedDocIdMultiValueIndexColumn<'a> {
|
||||
doc_id_mapping: &'a SegmentDocIdMapping,
|
||||
multi_value_length_readers: Vec<&'a MultiValueIndex>,
|
||||
min_value: u64,
|
||||
max_value: u64,
|
||||
num_vals: u32,
|
||||
}
|
||||
|
||||
impl<'a> RemappedDocIdMultiValueIndexColumn<'a> {
|
||||
pub(crate) fn new(
|
||||
segment_and_ff_readers: &'a [(&'a SegmentReader, &'a MultiValueIndex)],
|
||||
doc_id_mapping: &'a SegmentDocIdMapping,
|
||||
) -> Self {
|
||||
// We go through a complete first pass to compute the minimum and the
|
||||
// maximum value and initialize our Column.
|
||||
let mut num_vals = 0;
|
||||
let min_value = 0;
|
||||
let mut max_value = 0;
|
||||
let mut multi_value_length_readers = Vec::with_capacity(segment_and_ff_readers.len());
|
||||
for segment_and_ff_reader in segment_and_ff_readers {
|
||||
let segment_reader = segment_and_ff_reader.0;
|
||||
let multi_value_length_reader = segment_and_ff_reader.1;
|
||||
if !segment_reader.has_deletes() {
|
||||
max_value += multi_value_length_reader.total_num_vals() as u64;
|
||||
} else {
|
||||
for doc in segment_reader.doc_ids_alive() {
|
||||
max_value += multi_value_length_reader.num_vals_for_doc(doc) as u64;
|
||||
}
|
||||
}
|
||||
num_vals += segment_reader.num_docs();
|
||||
multi_value_length_readers.push(multi_value_length_reader);
|
||||
}
|
||||
// The value range is always get_val(doc)..get_val(doc + 1)
|
||||
num_vals += 1;
|
||||
Self {
|
||||
doc_id_mapping,
|
||||
multi_value_length_readers,
|
||||
min_value,
|
||||
max_value,
|
||||
num_vals,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Column for RemappedDocIdMultiValueIndexColumn<'a> {
|
||||
fn get_val(&self, _pos: u32) -> u64 {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
|
||||
let mut offset = 0;
|
||||
Box::new(
|
||||
std::iter::once(0).chain(self.doc_id_mapping.iter_old_doc_addrs().map(
|
||||
move |old_doc_addr| {
|
||||
let ff_reader =
|
||||
&self.multi_value_length_readers[old_doc_addr.segment_ord as usize];
|
||||
offset += ff_reader.num_vals_for_doc(old_doc_addr.doc_id);
|
||||
offset as u64
|
||||
},
|
||||
)),
|
||||
)
|
||||
}
|
||||
fn min_value(&self) -> u64 {
|
||||
self.min_value
|
||||
}
|
||||
|
||||
fn max_value(&self) -> u64 {
|
||||
self.max_value
|
||||
}
|
||||
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.num_vals
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,3 @@
|
||||
use itertools::Itertools;
|
||||
|
||||
use super::*;
|
||||
use crate::{Cardinality, ColumnarWriter, HasAssociatedColumnType, RowId};
|
||||
|
||||
@@ -14,7 +12,7 @@ fn make_columnar<T: Into<NumericalValue> + HasAssociatedColumnType + Copy>(
|
||||
}
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer
|
||||
.serialize(vals.len() as RowId, None, &mut buffer)
|
||||
.serialize(vals.len() as RowId, &mut buffer)
|
||||
.unwrap();
|
||||
ColumnarReader::open(buffer).unwrap()
|
||||
}
|
||||
@@ -26,7 +24,7 @@ fn test_column_coercion_to_u64() {
|
||||
// u64 type
|
||||
let columnar2 = make_columnar("numbers", &[u64::MAX]);
|
||||
let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
group_columns_for_merge(&[&columnar1, &columnar2], &[]).unwrap();
|
||||
group_columns_for_merge(&[&columnar1, &columnar2]).unwrap();
|
||||
assert_eq!(column_map.len(), 1);
|
||||
assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
|
||||
}
|
||||
@@ -36,7 +34,7 @@ fn test_column_no_coercion_if_all_the_same() {
|
||||
let columnar1 = make_columnar("numbers", &[1u64]);
|
||||
let columnar2 = make_columnar("numbers", &[2u64]);
|
||||
let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
group_columns_for_merge(&[&columnar1, &columnar2], &[]).unwrap();
|
||||
group_columns_for_merge(&[&columnar1, &columnar2]).unwrap();
|
||||
assert_eq!(column_map.len(), 1);
|
||||
assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
|
||||
}
|
||||
@@ -46,74 +44,17 @@ fn test_column_coercion_to_i64() {
|
||||
let columnar1 = make_columnar("numbers", &[-1i64]);
|
||||
let columnar2 = make_columnar("numbers", &[2u64]);
|
||||
let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
group_columns_for_merge(&[&columnar1, &columnar2], &[]).unwrap();
|
||||
group_columns_for_merge(&[&columnar1, &columnar2]).unwrap();
|
||||
assert_eq!(column_map.len(), 1);
|
||||
assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::I64)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_impossible_coercion_returns_an_error() {
|
||||
let columnar1 = make_columnar("numbers", &[u64::MAX]);
|
||||
let group_error =
|
||||
group_columns_for_merge(&[&columnar1], &[("numbers".to_string(), ColumnType::I64)])
|
||||
.map(|_| ())
|
||||
.unwrap_err();
|
||||
assert_eq!(group_error.kind(), io::ErrorKind::InvalidInput);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_group_columns_with_required_column() {
|
||||
let columnar1 = make_columnar("numbers", &[1i64]);
|
||||
let columnar2 = make_columnar("numbers", &[2u64]);
|
||||
let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
group_columns_for_merge(
|
||||
&[&columnar1, &columnar2],
|
||||
&[("numbers".to_string(), ColumnType::U64)],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(column_map.len(), 1);
|
||||
assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_group_columns_required_column_with_no_existing_columns() {
|
||||
let columnar1 = make_columnar("numbers", &[2u64]);
|
||||
let columnar2 = make_columnar("numbers", &[2u64]);
|
||||
let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
group_columns_for_merge(
|
||||
&[&columnar1, &columnar2],
|
||||
&[("required_col".to_string(), ColumnType::Str)],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(column_map.len(), 2);
|
||||
let columns = column_map
|
||||
.get(&("required_col".to_string(), ColumnType::Str))
|
||||
.unwrap();
|
||||
assert_eq!(columns.len(), 2);
|
||||
assert!(columns[0].is_none());
|
||||
assert!(columns[1].is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_group_columns_required_column_is_above_all_columns_have_the_same_type_rule() {
|
||||
let columnar1 = make_columnar("numbers", &[2i64]);
|
||||
let columnar2 = make_columnar("numbers", &[2i64]);
|
||||
let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
group_columns_for_merge(
|
||||
&[&columnar1, &columnar2],
|
||||
&[("numbers".to_string(), ColumnType::U64)],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(column_map.len(), 1);
|
||||
assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::U64)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_missing_column() {
|
||||
let columnar1 = make_columnar("numbers", &[-1i64]);
|
||||
let columnar2 = make_columnar("numbers2", &[2u64]);
|
||||
let column_map: BTreeMap<(String, ColumnType), Vec<Option<DynamicColumn>>> =
|
||||
group_columns_for_merge(&[&columnar1, &columnar2], &[]).unwrap();
|
||||
group_columns_for_merge(&[&columnar1, &columnar2]).unwrap();
|
||||
assert_eq!(column_map.len(), 2);
|
||||
assert!(column_map.contains_key(&("numbers".to_string(), ColumnType::I64)));
|
||||
{
|
||||
@@ -149,43 +90,16 @@ fn make_numerical_columnar_multiple_columns(
|
||||
.max()
|
||||
.unwrap_or(0u32);
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer
|
||||
.serialize(num_rows, None, &mut buffer)
|
||||
.unwrap();
|
||||
dataframe_writer.serialize(num_rows, &mut buffer).unwrap();
|
||||
ColumnarReader::open(buffer).unwrap()
|
||||
}
|
||||
|
||||
#[track_caller]
|
||||
fn make_byte_columnar_multiple_columns(
|
||||
columns: &[(&str, &[&[&[u8]]])],
|
||||
num_rows: u32,
|
||||
) -> ColumnarReader {
|
||||
let mut dataframe_writer = ColumnarWriter::default();
|
||||
for (column_name, column_values) in columns {
|
||||
assert_eq!(
|
||||
column_values.len(),
|
||||
num_rows as usize,
|
||||
"All columns must have `{num_rows}` rows"
|
||||
);
|
||||
for (row_id, vals) in column_values.iter().enumerate() {
|
||||
for val in vals.iter() {
|
||||
dataframe_writer.record_bytes(row_id as u32, column_name, val);
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer
|
||||
.serialize(num_rows, None, &mut buffer)
|
||||
.unwrap();
|
||||
ColumnarReader::open(buffer).unwrap()
|
||||
}
|
||||
|
||||
fn make_text_columnar_multiple_columns(columns: &[(&str, &[&[&str]])]) -> ColumnarReader {
|
||||
fn make_byte_columnar_multiple_columns(columns: &[(&str, &[&[&[u8]]])]) -> ColumnarReader {
|
||||
let mut dataframe_writer = ColumnarWriter::default();
|
||||
for (column_name, column_values) in columns {
|
||||
for (row_id, vals) in column_values.iter().enumerate() {
|
||||
for val in vals.iter() {
|
||||
dataframe_writer.record_str(row_id as u32, column_name, val);
|
||||
dataframe_writer.record_bytes(row_id as u32, column_name, *val);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -195,9 +109,26 @@ fn make_text_columnar_multiple_columns(columns: &[(&str, &[&[&str]])]) -> Column
|
||||
.max()
|
||||
.unwrap_or(0u32);
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer
|
||||
.serialize(num_rows, None, &mut buffer)
|
||||
.unwrap();
|
||||
dataframe_writer.serialize(num_rows, &mut buffer).unwrap();
|
||||
ColumnarReader::open(buffer).unwrap()
|
||||
}
|
||||
|
||||
fn make_text_columnar_multiple_columns(columns: &[(&str, &[&[&str]])]) -> ColumnarReader {
|
||||
let mut dataframe_writer = ColumnarWriter::default();
|
||||
for (column_name, column_values) in columns {
|
||||
for (row_id, vals) in column_values.iter().enumerate() {
|
||||
for val in vals.iter() {
|
||||
dataframe_writer.record_str(row_id as u32, column_name, *val);
|
||||
}
|
||||
}
|
||||
}
|
||||
let num_rows = columns
|
||||
.iter()
|
||||
.map(|(_, val_rows)| val_rows.len() as RowId)
|
||||
.max()
|
||||
.unwrap_or(0u32);
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(num_rows, &mut buffer).unwrap();
|
||||
ColumnarReader::open(buffer).unwrap()
|
||||
}
|
||||
|
||||
@@ -214,7 +145,6 @@ fn test_merge_columnar_numbers() {
|
||||
let stack_merge_order = StackMergeOrder::stack(columnars);
|
||||
crate::columnar::merge_columnar(
|
||||
columnars,
|
||||
&[],
|
||||
MergeRowOrder::Stack(stack_merge_order),
|
||||
&mut buffer,
|
||||
)
|
||||
@@ -240,7 +170,6 @@ fn test_merge_columnar_texts() {
|
||||
let stack_merge_order = StackMergeOrder::stack(columnars);
|
||||
crate::columnar::merge_columnar(
|
||||
columnars,
|
||||
&[],
|
||||
MergeRowOrder::Stack(stack_merge_order),
|
||||
&mut buffer,
|
||||
)
|
||||
@@ -251,8 +180,6 @@ fn test_merge_columnar_texts() {
|
||||
let cols = columnar_reader.read_columns("texts").unwrap();
|
||||
let dynamic_column = cols[0].open().unwrap();
|
||||
let DynamicColumn::Str(vals) = dynamic_column else { panic!() };
|
||||
assert_eq!(vals.ords().get_cardinality(), Cardinality::Optional);
|
||||
|
||||
let get_str_for_ord = |ord| {
|
||||
let mut out = String::new();
|
||||
vals.ord_to_str(ord, &mut out).unwrap();
|
||||
@@ -280,14 +207,13 @@ fn test_merge_columnar_texts() {
|
||||
|
||||
#[test]
|
||||
fn test_merge_columnar_byte() {
|
||||
let columnar1 = make_byte_columnar_multiple_columns(&[("bytes", &[&[b"bbbb"], &[b"baaa"]])], 2);
|
||||
let columnar2 = make_byte_columnar_multiple_columns(&[("bytes", &[&[], &[b"a"]])], 2);
|
||||
let columnar1 = make_byte_columnar_multiple_columns(&[("bytes", &[&[b"bbbb"], &[b"baaa"]])]);
|
||||
let columnar2 = make_byte_columnar_multiple_columns(&[("bytes", &[&[], &[b"a"]])]);
|
||||
let mut buffer = Vec::new();
|
||||
let columnars = &[&columnar1, &columnar2];
|
||||
let stack_merge_order = StackMergeOrder::stack(columnars);
|
||||
crate::columnar::merge_columnar(
|
||||
columnars,
|
||||
&[],
|
||||
MergeRowOrder::Stack(stack_merge_order),
|
||||
&mut buffer,
|
||||
)
|
||||
@@ -324,149 +250,3 @@ fn test_merge_columnar_byte() {
|
||||
assert_eq!(get_bytes_for_row(2), b"");
|
||||
assert_eq!(get_bytes_for_row(3), b"a");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_columnar_byte_with_missing() {
|
||||
let columnar1 = make_byte_columnar_multiple_columns(&[], 3);
|
||||
let columnar2 = make_byte_columnar_multiple_columns(&[("col", &[&[b"b"], &[]])], 2);
|
||||
let columnar3 = make_byte_columnar_multiple_columns(
|
||||
&[
|
||||
("col", &[&[], &[b"b"], &[b"a", b"b"]]),
|
||||
("col2", &[&[b"hello"], &[], &[b"a", b"b"]]),
|
||||
],
|
||||
3,
|
||||
);
|
||||
let mut buffer = Vec::new();
|
||||
let columnars = &[&columnar1, &columnar2, &columnar3];
|
||||
let stack_merge_order = StackMergeOrder::stack(columnars);
|
||||
crate::columnar::merge_columnar(
|
||||
columnars,
|
||||
&[],
|
||||
MergeRowOrder::Stack(stack_merge_order),
|
||||
&mut buffer,
|
||||
)
|
||||
.unwrap();
|
||||
let columnar_reader = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar_reader.num_rows(), 3 + 2 + 3);
|
||||
assert_eq!(columnar_reader.num_columns(), 2);
|
||||
let cols = columnar_reader.read_columns("col").unwrap();
|
||||
let dynamic_column = cols[0].open().unwrap();
|
||||
let DynamicColumn::Bytes(vals) = dynamic_column else { panic!() };
|
||||
let get_bytes_for_ord = |ord| {
|
||||
let mut out = Vec::new();
|
||||
vals.ord_to_bytes(ord, &mut out).unwrap();
|
||||
out
|
||||
};
|
||||
assert_eq!(vals.dictionary.num_terms(), 2);
|
||||
assert_eq!(get_bytes_for_ord(0), b"a");
|
||||
assert_eq!(get_bytes_for_ord(1), b"b");
|
||||
let get_bytes_for_row = |row_id| {
|
||||
let terms: Vec<Vec<u8>> = vals
|
||||
.term_ords(row_id)
|
||||
.map(|term_ord| {
|
||||
let mut out = Vec::new();
|
||||
vals.ord_to_bytes(term_ord, &mut out).unwrap();
|
||||
out
|
||||
})
|
||||
.collect();
|
||||
terms
|
||||
};
|
||||
assert!(get_bytes_for_row(0).is_empty());
|
||||
assert!(get_bytes_for_row(1).is_empty());
|
||||
assert!(get_bytes_for_row(2).is_empty());
|
||||
assert_eq!(get_bytes_for_row(3), vec![b"b".to_vec()]);
|
||||
assert!(get_bytes_for_row(4).is_empty());
|
||||
assert!(get_bytes_for_row(5).is_empty());
|
||||
assert_eq!(get_bytes_for_row(6), vec![b"b".to_vec()]);
|
||||
assert_eq!(get_bytes_for_row(7), vec![b"a".to_vec(), b"b".to_vec()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_columnar_different_types() {
|
||||
let columnar1 = make_text_columnar_multiple_columns(&[("mixed", &[&["a"]])]);
|
||||
let columnar2 = make_text_columnar_multiple_columns(&[("mixed", &[&[], &["b"]])]);
|
||||
let columnar3 = make_columnar("mixed", &[1i64]);
|
||||
let mut buffer = Vec::new();
|
||||
let columnars = &[&columnar1, &columnar2, &columnar3];
|
||||
let stack_merge_order = StackMergeOrder::stack(columnars);
|
||||
crate::columnar::merge_columnar(
|
||||
columnars,
|
||||
&[],
|
||||
MergeRowOrder::Stack(stack_merge_order),
|
||||
&mut buffer,
|
||||
)
|
||||
.unwrap();
|
||||
let columnar_reader = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar_reader.num_rows(), 4);
|
||||
assert_eq!(columnar_reader.num_columns(), 2);
|
||||
let cols = columnar_reader.read_columns("mixed").unwrap();
|
||||
|
||||
// numeric column
|
||||
let dynamic_column = cols[0].open().unwrap();
|
||||
let DynamicColumn::I64(vals) = dynamic_column else { panic!() };
|
||||
assert_eq!(vals.get_cardinality(), Cardinality::Optional);
|
||||
assert_eq!(vals.values_for_doc(0).collect_vec(), vec![]);
|
||||
assert_eq!(vals.values_for_doc(1).collect_vec(), vec![]);
|
||||
assert_eq!(vals.values_for_doc(2).collect_vec(), vec![]);
|
||||
assert_eq!(vals.values_for_doc(3).collect_vec(), vec![1]);
|
||||
assert_eq!(vals.values_for_doc(4).collect_vec(), vec![]);
|
||||
|
||||
// text column
|
||||
let dynamic_column = cols[1].open().unwrap();
|
||||
let DynamicColumn::Str(vals) = dynamic_column else { panic!() };
|
||||
assert_eq!(vals.ords().get_cardinality(), Cardinality::Optional);
|
||||
let get_str_for_ord = |ord| {
|
||||
let mut out = String::new();
|
||||
vals.ord_to_str(ord, &mut out).unwrap();
|
||||
out
|
||||
};
|
||||
|
||||
assert_eq!(vals.dictionary.num_terms(), 2);
|
||||
assert_eq!(get_str_for_ord(0), "a");
|
||||
assert_eq!(get_str_for_ord(1), "b");
|
||||
|
||||
let get_str_for_row = |row_id| {
|
||||
let term_ords: Vec<String> = vals
|
||||
.term_ords(row_id)
|
||||
.map(|el| {
|
||||
let mut out = String::new();
|
||||
vals.ord_to_str(el, &mut out).unwrap();
|
||||
out
|
||||
})
|
||||
.collect();
|
||||
term_ords
|
||||
};
|
||||
|
||||
assert_eq!(get_str_for_row(0), vec!["a".to_string()]);
|
||||
assert_eq!(get_str_for_row(1), Vec::<String>::new());
|
||||
assert_eq!(get_str_for_row(2), vec!["b".to_string()]);
|
||||
assert_eq!(get_str_for_row(3), Vec::<String>::new());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_columnar_different_empty_cardinality() {
|
||||
let columnar1 = make_text_columnar_multiple_columns(&[("mixed", &[&["a"]])]);
|
||||
let columnar2 = make_columnar("mixed", &[1i64]);
|
||||
let mut buffer = Vec::new();
|
||||
let columnars = &[&columnar1, &columnar2];
|
||||
let stack_merge_order = StackMergeOrder::stack(columnars);
|
||||
crate::columnar::merge_columnar(
|
||||
columnars,
|
||||
&[],
|
||||
MergeRowOrder::Stack(stack_merge_order),
|
||||
&mut buffer,
|
||||
)
|
||||
.unwrap();
|
||||
let columnar_reader = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar_reader.num_rows(), 2);
|
||||
assert_eq!(columnar_reader.num_columns(), 2);
|
||||
let cols = columnar_reader.read_columns("mixed").unwrap();
|
||||
|
||||
// numeric column
|
||||
let dynamic_column = cols[0].open().unwrap();
|
||||
assert_eq!(dynamic_column.get_cardinality(), Cardinality::Optional);
|
||||
|
||||
// text column
|
||||
let dynamic_column = cols[1].open().unwrap();
|
||||
assert_eq!(dynamic_column.get_cardinality(), Cardinality::Optional);
|
||||
}
|
||||
|
||||
1
columnar/src/columnar/merge_index.rs
Normal file
1
columnar/src/columnar/merge_index.rs
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
@@ -1,12 +1,11 @@
|
||||
mod column_type;
|
||||
mod format_version;
|
||||
mod merge;
|
||||
mod merge_index;
|
||||
mod reader;
|
||||
mod writer;
|
||||
|
||||
pub use column_type::{ColumnType, HasAssociatedColumnType};
|
||||
#[cfg(test)]
|
||||
pub(crate) use merge::ColumnTypeCategory;
|
||||
pub use merge::{merge_columnar, MergeRowOrder, ShuffleMergeOrder, StackMergeOrder};
|
||||
pub use merge::{merge_columnar, MergeRowOrder, StackMergeOrder};
|
||||
pub use reader::ColumnarReader;
|
||||
pub use writer::ColumnarWriter;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::{fmt, io, mem};
|
||||
use std::{io, mem};
|
||||
|
||||
use common::file_slice::FileSlice;
|
||||
use common::BinarySerializable;
|
||||
@@ -21,58 +21,6 @@ pub struct ColumnarReader {
|
||||
num_rows: RowId,
|
||||
}
|
||||
|
||||
impl fmt::Debug for ColumnarReader {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let num_rows = self.num_rows();
|
||||
let columns = self.list_columns().unwrap();
|
||||
let num_cols = columns.len();
|
||||
let mut debug_struct = f.debug_struct("Columnar");
|
||||
debug_struct
|
||||
.field("num_rows", &num_rows)
|
||||
.field("num_cols", &num_cols);
|
||||
for (col_name, dynamic_column_handle) in columns.into_iter().take(5) {
|
||||
let col = dynamic_column_handle.open().unwrap();
|
||||
if col.num_values() > 10 {
|
||||
debug_struct.field(&col_name, &"..");
|
||||
} else {
|
||||
debug_struct.field(&col_name, &col);
|
||||
}
|
||||
}
|
||||
if num_cols > 5 {
|
||||
debug_struct.finish_non_exhaustive()?;
|
||||
} else {
|
||||
debug_struct.finish()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Functions by both the async/sync code listing columns.
|
||||
/// It takes a stream from the column sstable and return the list of
|
||||
/// `DynamicColumn` available in it.
|
||||
fn read_all_columns_in_stream(
|
||||
mut stream: sstable::Streamer<'_, RangeSSTable>,
|
||||
column_data: &FileSlice,
|
||||
) -> io::Result<Vec<DynamicColumnHandle>> {
|
||||
let mut results = Vec::new();
|
||||
while stream.advance() {
|
||||
let key_bytes: &[u8] = stream.key();
|
||||
let Some(column_code) = key_bytes.last().copied() else {
|
||||
return Err(io_invalid_data("Empty column name.".to_string()));
|
||||
};
|
||||
let column_type = ColumnType::try_from_code(column_code)
|
||||
.map_err(|_| io_invalid_data(format!("Unknown column code `{column_code}`")))?;
|
||||
let range = stream.value();
|
||||
let file_slice = column_data.slice(range.start as usize..range.end as usize);
|
||||
let dynamic_column_handle = DynamicColumnHandle {
|
||||
file_slice,
|
||||
column_type,
|
||||
};
|
||||
results.push(dynamic_column_handle);
|
||||
}
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
impl ColumnarReader {
|
||||
/// Opens a new Columnar file.
|
||||
pub fn open<F>(file_slice: F) -> io::Result<ColumnarReader>
|
||||
@@ -128,7 +76,11 @@ impl ColumnarReader {
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
fn stream_for_column_range(&self, column_name: &str) -> sstable::StreamerBuilder<RangeSSTable> {
|
||||
/// Get all columns for the given column name.
|
||||
///
|
||||
/// There can be more than one column associated to a given column name, provided they have
|
||||
/// different types.
|
||||
pub fn read_columns(&self, column_name: &str) -> io::Result<Vec<DynamicColumnHandle>> {
|
||||
// Each column is a associated to a given `column_key`,
|
||||
// that starts by `column_name\0column_header`.
|
||||
//
|
||||
@@ -137,35 +89,36 @@ impl ColumnarReader {
|
||||
//
|
||||
// This is in turn equivalent to searching for the range
|
||||
// `[column_name,\0`..column_name\1)`.
|
||||
// TODO can we get some more generic `prefix(..)` logic in the dictionary.
|
||||
|
||||
// TODO can we get some more generic `prefix(..)` logic in the dictioanry.
|
||||
let mut start_key = column_name.to_string();
|
||||
start_key.push('\0');
|
||||
let mut end_key = column_name.to_string();
|
||||
end_key.push(1u8 as char);
|
||||
self.column_dictionary
|
||||
let mut stream = self
|
||||
.column_dictionary
|
||||
.range()
|
||||
.ge(start_key.as_bytes())
|
||||
.lt(end_key.as_bytes())
|
||||
}
|
||||
|
||||
pub async fn read_columns_async(
|
||||
&self,
|
||||
column_name: &str,
|
||||
) -> io::Result<Vec<DynamicColumnHandle>> {
|
||||
let stream = self
|
||||
.stream_for_column_range(column_name)
|
||||
.into_stream_async()
|
||||
.await?;
|
||||
read_all_columns_in_stream(stream, &self.column_data)
|
||||
}
|
||||
|
||||
/// Get all columns for the given column name.
|
||||
///
|
||||
/// There can be more than one column associated to a given column name, provided they have
|
||||
/// different types.
|
||||
pub fn read_columns(&self, column_name: &str) -> io::Result<Vec<DynamicColumnHandle>> {
|
||||
let stream = self.stream_for_column_range(column_name).into_stream()?;
|
||||
read_all_columns_in_stream(stream, &self.column_data)
|
||||
.into_stream()?;
|
||||
let mut results = Vec::new();
|
||||
while stream.advance() {
|
||||
let key_bytes: &[u8] = stream.key();
|
||||
assert!(key_bytes.starts_with(start_key.as_bytes()));
|
||||
let column_code: u8 = key_bytes.last().cloned().unwrap();
|
||||
let column_type = ColumnType::try_from_code(column_code)
|
||||
.map_err(|_| io_invalid_data(format!("Unknown column code `{column_code}`")))?;
|
||||
let range = stream.value().clone();
|
||||
let file_slice = self
|
||||
.column_data
|
||||
.slice(range.start as usize..range.end as usize);
|
||||
let dynamic_column_handle = DynamicColumnHandle {
|
||||
file_slice,
|
||||
column_type,
|
||||
};
|
||||
results.push(dynamic_column_handle);
|
||||
}
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Return the number of columns in the columnar.
|
||||
@@ -184,7 +137,7 @@ mod tests {
|
||||
columnar_writer.record_column_type("col1", ColumnType::Str, false);
|
||||
columnar_writer.record_column_type("col2", ColumnType::U64, false);
|
||||
let mut buffer = Vec::new();
|
||||
columnar_writer.serialize(1, None, &mut buffer).unwrap();
|
||||
columnar_writer.serialize(1, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
let columns = columnar.list_columns().unwrap();
|
||||
assert_eq!(columns.len(), 2);
|
||||
@@ -200,7 +153,7 @@ mod tests {
|
||||
columnar_writer.record_column_type("count", ColumnType::U64, false);
|
||||
columnar_writer.record_numerical(1, "count", 1u64);
|
||||
let mut buffer = Vec::new();
|
||||
columnar_writer.serialize(2, None, &mut buffer).unwrap();
|
||||
columnar_writer.serialize(2, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
let columns = columnar.list_columns().unwrap();
|
||||
assert_eq!(columns.len(), 1);
|
||||
@@ -209,7 +162,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "Input type forbidden")]
|
||||
#[should_panic(expect = "Input type forbidden")]
|
||||
fn test_list_columns_strict_typing_panics_on_wrong_types() {
|
||||
let mut columnar_writer = ColumnarWriter::default();
|
||||
columnar_writer.record_column_type("count", ColumnType::U64, false);
|
||||
|
||||
@@ -310,7 +310,7 @@ mod tests {
|
||||
buffer.extend_from_slice(b"234234");
|
||||
let mut bytes = &buffer[..];
|
||||
let serdeser_symbol = ColumnOperation::deserialize(&mut bytes).unwrap();
|
||||
assert_eq!(bytes.len() + buf.as_ref().len(), buffer.len());
|
||||
assert_eq!(bytes.len() + buf.as_ref().len() as usize, buffer.len());
|
||||
assert_eq!(column_op, serdeser_symbol);
|
||||
}
|
||||
|
||||
@@ -341,7 +341,7 @@ mod tests {
|
||||
fn test_column_operation_unordered_aux(val: u32, expected_len: usize) {
|
||||
let column_op = ColumnOperation::Value(UnorderedId(val));
|
||||
let minibuf = column_op.serialize();
|
||||
assert_eq!({ minibuf.as_ref().len() }, expected_len);
|
||||
assert_eq!(minibuf.as_ref().len() as usize, expected_len);
|
||||
let mut buf = minibuf.as_ref().to_vec();
|
||||
buf.extend_from_slice(&[2, 2, 2, 2, 2, 2]);
|
||||
let mut cursor = &buf[..];
|
||||
|
||||
@@ -41,31 +41,10 @@ impl ColumnWriter {
|
||||
pub(super) fn operation_iterator<'a, V: SymbolValue>(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
old_to_new_ids_opt: Option<&[RowId]>,
|
||||
buffer: &'a mut Vec<u8>,
|
||||
) -> impl Iterator<Item = ColumnOperation<V>> + 'a {
|
||||
buffer.clear();
|
||||
self.values.read_to_end(arena, buffer);
|
||||
if let Some(old_to_new_ids) = old_to_new_ids_opt {
|
||||
// TODO avoid the extra deserialization / serialization.
|
||||
let mut sorted_ops: Vec<(RowId, ColumnOperation<V>)> = Vec::new();
|
||||
let mut new_doc = 0u32;
|
||||
let mut cursor = &buffer[..];
|
||||
for op in std::iter::from_fn(|| ColumnOperation::<V>::deserialize(&mut cursor)) {
|
||||
if let ColumnOperation::NewDoc(doc) = &op {
|
||||
new_doc = old_to_new_ids[*doc as usize];
|
||||
sorted_ops.push((new_doc, ColumnOperation::NewDoc(new_doc)));
|
||||
} else {
|
||||
sorted_ops.push((new_doc, op));
|
||||
}
|
||||
}
|
||||
// stable sort is crucial here.
|
||||
sorted_ops.sort_by_key(|(new_doc_id, _)| *new_doc_id);
|
||||
buffer.clear();
|
||||
for (_, op) in sorted_ops {
|
||||
buffer.extend_from_slice(op.serialize().as_ref());
|
||||
}
|
||||
}
|
||||
let mut cursor: &[u8] = &buffer[..];
|
||||
std::iter::from_fn(move || ColumnOperation::deserialize(&mut cursor))
|
||||
}
|
||||
@@ -210,12 +189,10 @@ impl CompatibleNumericalTypes {
|
||||
}
|
||||
|
||||
impl NumericalColumnWriter {
|
||||
pub fn numerical_type(&self) -> NumericalType {
|
||||
self.compatible_numerical_types.to_numerical_type()
|
||||
}
|
||||
|
||||
pub fn cardinality(&self, num_docs: RowId) -> Cardinality {
|
||||
self.column_writer.get_cardinality(num_docs)
|
||||
pub fn column_type_and_cardinality(&self, num_docs: RowId) -> (NumericalType, Cardinality) {
|
||||
let numerical_type = self.compatible_numerical_types.to_numerical_type();
|
||||
let cardinality = self.column_writer.get_cardinality(num_docs);
|
||||
(numerical_type, cardinality)
|
||||
}
|
||||
|
||||
pub fn record_numerical_value(
|
||||
@@ -231,11 +208,9 @@ impl NumericalColumnWriter {
|
||||
pub(super) fn operation_iterator<'a>(
|
||||
self,
|
||||
arena: &MemoryArena,
|
||||
old_to_new_ids: Option<&[RowId]>,
|
||||
buffer: &'a mut Vec<u8>,
|
||||
) -> impl Iterator<Item = ColumnOperation<NumericalValue>> + 'a {
|
||||
self.column_writer
|
||||
.operation_iterator(arena, old_to_new_ids, buffer)
|
||||
self.column_writer.operation_iterator(arena, buffer)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -276,11 +251,9 @@ impl StrOrBytesColumnWriter {
|
||||
pub(super) fn operation_iterator<'a>(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
old_to_new_ids: Option<&[RowId]>,
|
||||
byte_buffer: &'a mut Vec<u8>,
|
||||
) -> impl Iterator<Item = ColumnOperation<UnorderedId>> + 'a {
|
||||
self.column_writer
|
||||
.operation_iterator(arena, old_to_new_ids, byte_buffer)
|
||||
self.column_writer.operation_iterator(arena, byte_buffer)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ use crate::column_index::SerializableColumnIndex;
|
||||
use crate::column_values::{
|
||||
ColumnValues, MonotonicallyMappableToU128, MonotonicallyMappableToU64, VecColumn,
|
||||
};
|
||||
use crate::columnar::column_type::ColumnType;
|
||||
use crate::columnar::column_type::{ColumnType, ColumnTypeCategory};
|
||||
use crate::columnar::writer::column_writers::{
|
||||
ColumnWriter, NumericalColumnWriter, StrOrBytesColumnWriter,
|
||||
};
|
||||
@@ -45,9 +45,8 @@ struct SpareBuffers {
|
||||
/// columnar_writer.record_str(1u32 /* doc id */, "product_name", "Apple");
|
||||
/// columnar_writer.record_numerical(0u32 /* doc id */, "price", 10.5f64); //< uh oh we ended up mixing integer and floats.
|
||||
/// let mut wrt: Vec<u8> = Vec::new();
|
||||
/// columnar_writer.serialize(2u32, None, &mut wrt).unwrap();
|
||||
/// columnar_writer.serialize(2u32, &mut wrt).unwrap();
|
||||
/// ```
|
||||
#[derive(Default)]
|
||||
pub struct ColumnarWriter {
|
||||
numerical_field_hash_map: ArenaHashMap,
|
||||
datetime_field_hash_map: ArenaHashMap,
|
||||
@@ -61,6 +60,22 @@ pub struct ColumnarWriter {
|
||||
buffers: SpareBuffers,
|
||||
}
|
||||
|
||||
impl Default for ColumnarWriter {
|
||||
fn default() -> Self {
|
||||
ColumnarWriter {
|
||||
numerical_field_hash_map: ArenaHashMap::new(10_000),
|
||||
bool_field_hash_map: ArenaHashMap::new(10_000),
|
||||
ip_addr_field_hash_map: ArenaHashMap::new(10_000),
|
||||
bytes_field_hash_map: ArenaHashMap::new(10_000),
|
||||
str_field_hash_map: ArenaHashMap::new(10_000),
|
||||
datetime_field_hash_map: ArenaHashMap::new(10_000),
|
||||
dictionaries: Vec::new(),
|
||||
arena: MemoryArena::default(),
|
||||
buffers: SpareBuffers::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn mutate_or_create_column<V, TMutator>(
|
||||
arena_hash_map: &mut ArenaHashMap,
|
||||
@@ -89,48 +104,6 @@ impl ColumnarWriter {
|
||||
+ self.datetime_field_hash_map.mem_usage()
|
||||
}
|
||||
|
||||
/// Returns the list of doc ids from 0..num_docs sorted by the `sort_field`
|
||||
/// column.
|
||||
///
|
||||
/// If the column is multivalued, use the first value for scoring.
|
||||
/// If no value is associated to a specific row, the document is assigned
|
||||
/// the lowest possible score.
|
||||
///
|
||||
/// The sort applied is stable.
|
||||
pub fn sort_order(&self, sort_field: &str, num_docs: RowId, reversed: bool) -> Vec<u32> {
|
||||
let Some(numerical_col_writer) =
|
||||
self.numerical_field_hash_map.get::<NumericalColumnWriter>(sort_field.as_bytes()) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut symbols_buffer = Vec::new();
|
||||
let mut values = Vec::new();
|
||||
let mut last_doc_opt: Option<RowId> = None;
|
||||
for op in numerical_col_writer.operation_iterator(&self.arena, None, &mut symbols_buffer) {
|
||||
match op {
|
||||
ColumnOperation::NewDoc(doc) => {
|
||||
last_doc_opt = Some(doc);
|
||||
}
|
||||
ColumnOperation::Value(numerical_value) => {
|
||||
if let Some(last_doc) = last_doc_opt {
|
||||
let score: f32 = f64::coerce(numerical_value) as f32;
|
||||
values.push((score, last_doc));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for doc in values.len() as u32..num_docs {
|
||||
values.push((0.0f32, doc));
|
||||
}
|
||||
values.sort_by(|(left_score, _), (right_score, _)| {
|
||||
if reversed {
|
||||
right_score.partial_cmp(left_score).unwrap()
|
||||
} else {
|
||||
left_score.partial_cmp(right_score).unwrap()
|
||||
}
|
||||
});
|
||||
values.into_iter().map(|(_score, doc)| doc).collect()
|
||||
}
|
||||
|
||||
/// Records a column type. This is useful to bypass the coercion process,
|
||||
/// makes sure the empty is present in the resulting columnar, or set
|
||||
/// the `sort_values_within_row`.
|
||||
@@ -251,15 +224,11 @@ impl ColumnarWriter {
|
||||
});
|
||||
}
|
||||
|
||||
pub fn record_datetime(&mut self, doc: RowId, column_name: &str, datetime: common::DateTime) {
|
||||
pub fn record_datetime(&mut self, doc: RowId, column_name: &str, datetime: crate::DateTime) {
|
||||
let (hash_map, arena) = (&mut self.datetime_field_hash_map, &mut self.arena);
|
||||
mutate_or_create_column(hash_map, column_name, |column_opt: Option<ColumnWriter>| {
|
||||
let mut column: ColumnWriter = column_opt.unwrap_or_default();
|
||||
column.record(
|
||||
doc,
|
||||
NumericalValue::I64(datetime.into_timestamp_micros()),
|
||||
arena,
|
||||
);
|
||||
column.record(doc, NumericalValue::I64(datetime.timestamp_micros), arena);
|
||||
column
|
||||
});
|
||||
}
|
||||
@@ -309,47 +278,37 @@ impl ColumnarWriter {
|
||||
},
|
||||
);
|
||||
}
|
||||
pub fn serialize(
|
||||
&mut self,
|
||||
num_docs: RowId,
|
||||
old_to_new_row_ids: Option<&[RowId]>,
|
||||
wrt: &mut dyn io::Write,
|
||||
) -> io::Result<()> {
|
||||
pub fn serialize(&mut self, num_docs: RowId, wrt: &mut dyn io::Write) -> io::Result<()> {
|
||||
let mut serializer = ColumnarSerializer::new(wrt);
|
||||
let mut columns: Vec<(&[u8], ColumnType, Addr)> = self
|
||||
let mut columns: Vec<(&[u8], ColumnTypeCategory, Addr)> = self
|
||||
.numerical_field_hash_map
|
||||
.iter()
|
||||
.map(|(column_name, addr, _)| {
|
||||
let numerical_column_writer: NumericalColumnWriter =
|
||||
self.numerical_field_hash_map.read(addr);
|
||||
let column_type = numerical_column_writer.numerical_type().into();
|
||||
(column_name, column_type, addr)
|
||||
})
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnTypeCategory::Numerical, addr))
|
||||
.collect();
|
||||
columns.extend(
|
||||
self.bytes_field_hash_map
|
||||
.iter()
|
||||
.map(|(term, addr, _)| (term, ColumnType::Bytes, addr)),
|
||||
.map(|(term, addr, _)| (term, ColumnTypeCategory::Bytes, addr)),
|
||||
);
|
||||
columns.extend(
|
||||
self.str_field_hash_map
|
||||
.iter()
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnType::Str, addr)),
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnTypeCategory::Str, addr)),
|
||||
);
|
||||
columns.extend(
|
||||
self.bool_field_hash_map
|
||||
.iter()
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnType::Bool, addr)),
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnTypeCategory::Bool, addr)),
|
||||
);
|
||||
columns.extend(
|
||||
self.ip_addr_field_hash_map
|
||||
.iter()
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnType::IpAddr, addr)),
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnTypeCategory::IpAddr, addr)),
|
||||
);
|
||||
columns.extend(
|
||||
self.datetime_field_hash_map
|
||||
.iter()
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnType::DateTime, addr)),
|
||||
.map(|(column_name, addr, _)| (column_name, ColumnTypeCategory::DateTime, addr)),
|
||||
);
|
||||
columns.sort_unstable_by_key(|(column_name, col_type, _)| (*column_name, *col_type));
|
||||
|
||||
@@ -357,24 +316,20 @@ impl ColumnarWriter {
|
||||
let mut symbol_byte_buffer: Vec<u8> = Vec::new();
|
||||
for (column_name, column_type, addr) in columns {
|
||||
match column_type {
|
||||
ColumnType::Bool => {
|
||||
ColumnTypeCategory::Bool => {
|
||||
let column_writer: ColumnWriter = self.bool_field_hash_map.read(addr);
|
||||
let cardinality = column_writer.get_cardinality(num_docs);
|
||||
let mut column_serializer =
|
||||
serializer.serialize_column(column_name, column_type);
|
||||
serializer.serialize_column(column_name, ColumnType::Bool);
|
||||
serialize_bool_column(
|
||||
cardinality,
|
||||
num_docs,
|
||||
column_writer.operation_iterator(
|
||||
arena,
|
||||
old_to_new_row_ids,
|
||||
&mut symbol_byte_buffer,
|
||||
),
|
||||
column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
}
|
||||
ColumnType::IpAddr => {
|
||||
ColumnTypeCategory::IpAddr => {
|
||||
let column_writer: ColumnWriter = self.ip_addr_field_hash_map.read(addr);
|
||||
let cardinality = column_writer.get_cardinality(num_docs);
|
||||
let mut column_serializer =
|
||||
@@ -382,64 +337,50 @@ impl ColumnarWriter {
|
||||
serialize_ip_addr_column(
|
||||
cardinality,
|
||||
num_docs,
|
||||
column_writer.operation_iterator(
|
||||
arena,
|
||||
old_to_new_row_ids,
|
||||
&mut symbol_byte_buffer,
|
||||
),
|
||||
column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
}
|
||||
ColumnType::Bytes | ColumnType::Str => {
|
||||
let str_or_bytes_column_writer: StrOrBytesColumnWriter =
|
||||
if column_type == ColumnType::Bytes {
|
||||
self.bytes_field_hash_map.read(addr)
|
||||
ColumnTypeCategory::Bytes | ColumnTypeCategory::Str => {
|
||||
let (column_type, str_column_writer): (ColumnType, StrOrBytesColumnWriter) =
|
||||
if column_type == ColumnTypeCategory::Bytes {
|
||||
(ColumnType::Bytes, self.bytes_field_hash_map.read(addr))
|
||||
} else {
|
||||
self.str_field_hash_map.read(addr)
|
||||
(ColumnType::Str, self.str_field_hash_map.read(addr))
|
||||
};
|
||||
let dictionary_builder =
|
||||
&dictionaries[str_or_bytes_column_writer.dictionary_id as usize];
|
||||
let cardinality = str_or_bytes_column_writer
|
||||
.column_writer
|
||||
.get_cardinality(num_docs);
|
||||
&dictionaries[str_column_writer.dictionary_id as usize];
|
||||
let cardinality = str_column_writer.column_writer.get_cardinality(num_docs);
|
||||
let mut column_serializer =
|
||||
serializer.serialize_column(column_name, column_type);
|
||||
serialize_bytes_or_str_column(
|
||||
cardinality,
|
||||
num_docs,
|
||||
str_or_bytes_column_writer.sort_values_within_row,
|
||||
str_column_writer.sort_values_within_row,
|
||||
dictionary_builder,
|
||||
str_or_bytes_column_writer.operation_iterator(
|
||||
arena,
|
||||
old_to_new_row_ids,
|
||||
&mut symbol_byte_buffer,
|
||||
),
|
||||
str_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
}
|
||||
ColumnType::F64 | ColumnType::I64 | ColumnType::U64 => {
|
||||
ColumnTypeCategory::Numerical => {
|
||||
let numerical_column_writer: NumericalColumnWriter =
|
||||
self.numerical_field_hash_map.read(addr);
|
||||
let cardinality = numerical_column_writer.cardinality(num_docs);
|
||||
let (numerical_type, cardinality) =
|
||||
numerical_column_writer.column_type_and_cardinality(num_docs);
|
||||
let mut column_serializer =
|
||||
serializer.serialize_column(column_name, column_type);
|
||||
let numerical_type = column_type.numerical_type().unwrap();
|
||||
serializer.serialize_column(column_name, ColumnType::from(numerical_type));
|
||||
serialize_numerical_column(
|
||||
cardinality,
|
||||
num_docs,
|
||||
numerical_type,
|
||||
numerical_column_writer.operation_iterator(
|
||||
arena,
|
||||
old_to_new_row_ids,
|
||||
&mut symbol_byte_buffer,
|
||||
),
|
||||
numerical_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
}
|
||||
ColumnType::DateTime => {
|
||||
ColumnTypeCategory::DateTime => {
|
||||
let column_writer: ColumnWriter = self.datetime_field_hash_map.read(addr);
|
||||
let cardinality = column_writer.get_cardinality(num_docs);
|
||||
let mut column_serializer =
|
||||
@@ -448,11 +389,7 @@ impl ColumnarWriter {
|
||||
cardinality,
|
||||
num_docs,
|
||||
NumericalType::I64,
|
||||
column_writer.operation_iterator(
|
||||
arena,
|
||||
old_to_new_row_ids,
|
||||
&mut symbol_byte_buffer,
|
||||
),
|
||||
column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
@@ -651,12 +588,13 @@ where
|
||||
crate::column::serialize_column_mappable_to_u128(
|
||||
serializable_column_index,
|
||||
&&values[..],
|
||||
values.len() as u32,
|
||||
&mut wrt,
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn sort_values_within_row_in_place(multivalued_index: &[RowId], values: &mut [u64]) {
|
||||
fn sort_values_within_row_in_place(multivalued_index: &[RowId], values: &mut Vec<u64>) {
|
||||
let mut start_index: usize = 0;
|
||||
for end_index in multivalued_index.iter().copied() {
|
||||
let end_index = end_index as usize;
|
||||
@@ -761,7 +699,7 @@ mod tests {
|
||||
assert_eq!(column_writer.get_cardinality(3), Cardinality::Full);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&arena, None, &mut buffer)
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 6);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
|
||||
@@ -790,7 +728,7 @@ mod tests {
|
||||
assert_eq!(column_writer.get_cardinality(3), Cardinality::Optional);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&arena, None, &mut buffer)
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 4);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(1u32)));
|
||||
@@ -813,7 +751,7 @@ mod tests {
|
||||
assert_eq!(column_writer.get_cardinality(2), Cardinality::Optional);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&arena, None, &mut buffer)
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 2);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
|
||||
@@ -832,7 +770,7 @@ mod tests {
|
||||
assert_eq!(column_writer.get_cardinality(1), Cardinality::Multivalued);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&arena, None, &mut buffer)
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 3);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
|
||||
|
||||
@@ -29,7 +29,7 @@ pub struct OptionalIndexBuilder {
|
||||
}
|
||||
|
||||
impl OptionalIndexBuilder {
|
||||
pub fn finish(&mut self, num_rows: RowId) -> impl Iterable<RowId> + '_ {
|
||||
pub fn finish<'a>(&'a mut self, num_rows: RowId) -> impl Iterable<RowId> + 'a {
|
||||
debug_assert!(self
|
||||
.docs
|
||||
.last()
|
||||
@@ -150,7 +150,11 @@ mod tests {
|
||||
multivalued_value_index_builder.record_row(2u32);
|
||||
multivalued_value_index_builder.record_value();
|
||||
assert_eq!(
|
||||
multivalued_value_index_builder.finish(4u32).to_vec(),
|
||||
multivalued_value_index_builder
|
||||
.finish(4u32)
|
||||
.iter()
|
||||
.copied()
|
||||
.collect::<Vec<u32>>(),
|
||||
vec![0, 0, 2, 3, 3]
|
||||
);
|
||||
multivalued_value_index_builder.reset();
|
||||
@@ -158,7 +162,11 @@ mod tests {
|
||||
multivalued_value_index_builder.record_value();
|
||||
multivalued_value_index_builder.record_value();
|
||||
assert_eq!(
|
||||
multivalued_value_index_builder.finish(4u32).to_vec(),
|
||||
multivalued_value_index_builder
|
||||
.finish(4u32)
|
||||
.iter()
|
||||
.copied()
|
||||
.collect::<Vec<u32>>(),
|
||||
vec![0, 0, 0, 2, 2]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
use std::io;
|
||||
use std::net::Ipv6Addr;
|
||||
use std::sync::Arc;
|
||||
use std::{fmt, io};
|
||||
|
||||
use common::file_slice::FileSlice;
|
||||
use common::{ByteCount, DateTime, HasLen, OwnedBytes};
|
||||
use common::{HasLen, OwnedBytes};
|
||||
|
||||
use crate::column::{BytesColumn, Column, StrColumn};
|
||||
use crate::column_values::{monotonic_map_column, StrictlyMonotonicFn};
|
||||
use crate::columnar::ColumnType;
|
||||
use crate::{Cardinality, ColumnIndex, NumericalType};
|
||||
use crate::{Cardinality, DateTime, NumericalType};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum DynamicColumn {
|
||||
@@ -22,54 +22,19 @@ pub enum DynamicColumn {
|
||||
Str(StrColumn),
|
||||
}
|
||||
|
||||
impl fmt::Debug for DynamicColumn {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
write!(f, "[{} {} |", self.get_cardinality(), self.column_type())?;
|
||||
match self {
|
||||
DynamicColumn::Bool(col) => write!(f, " {:?}", col)?,
|
||||
DynamicColumn::I64(col) => write!(f, " {:?}", col)?,
|
||||
DynamicColumn::U64(col) => write!(f, " {:?}", col)?,
|
||||
DynamicColumn::F64(col) => write!(f, "{:?}", col)?,
|
||||
DynamicColumn::IpAddr(col) => write!(f, "{:?}", col)?,
|
||||
DynamicColumn::DateTime(col) => write!(f, "{:?}", col)?,
|
||||
DynamicColumn::Bytes(col) => write!(f, "{:?}", col)?,
|
||||
DynamicColumn::Str(col) => write!(f, "{:?}", col)?,
|
||||
}
|
||||
write!(f, "]")
|
||||
}
|
||||
}
|
||||
|
||||
impl DynamicColumn {
|
||||
pub fn column_index(&self) -> &ColumnIndex {
|
||||
match self {
|
||||
DynamicColumn::Bool(c) => &c.index,
|
||||
DynamicColumn::I64(c) => &c.index,
|
||||
DynamicColumn::U64(c) => &c.index,
|
||||
DynamicColumn::F64(c) => &c.index,
|
||||
DynamicColumn::IpAddr(c) => &c.index,
|
||||
DynamicColumn::DateTime(c) => &c.index,
|
||||
DynamicColumn::Bytes(c) => &c.ords().index,
|
||||
DynamicColumn::Str(c) => &c.ords().index,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_cardinality(&self) -> Cardinality {
|
||||
self.column_index().get_cardinality()
|
||||
}
|
||||
|
||||
pub fn num_values(&self) -> u32 {
|
||||
match self {
|
||||
DynamicColumn::Bool(c) => c.values.num_vals(),
|
||||
DynamicColumn::I64(c) => c.values.num_vals(),
|
||||
DynamicColumn::U64(c) => c.values.num_vals(),
|
||||
DynamicColumn::F64(c) => c.values.num_vals(),
|
||||
DynamicColumn::IpAddr(c) => c.values.num_vals(),
|
||||
DynamicColumn::DateTime(c) => c.values.num_vals(),
|
||||
DynamicColumn::Bytes(c) => c.ords().values.num_vals(),
|
||||
DynamicColumn::Str(c) => c.ords().values.num_vals(),
|
||||
DynamicColumn::Bool(c) => c.get_cardinality(),
|
||||
DynamicColumn::I64(c) => c.get_cardinality(),
|
||||
DynamicColumn::U64(c) => c.get_cardinality(),
|
||||
DynamicColumn::F64(c) => c.get_cardinality(),
|
||||
DynamicColumn::IpAddr(c) => c.get_cardinality(),
|
||||
DynamicColumn::DateTime(c) => c.get_cardinality(),
|
||||
DynamicColumn::Bytes(c) => c.ords().get_cardinality(),
|
||||
DynamicColumn::Str(c) => c.ords().get_cardinality(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn column_type(&self) -> ColumnType {
|
||||
match self {
|
||||
DynamicColumn::Bool(_) => ColumnType::Bool,
|
||||
@@ -108,11 +73,11 @@ impl DynamicColumn {
|
||||
fn coerce_to_f64(self) -> Option<DynamicColumn> {
|
||||
match self {
|
||||
DynamicColumn::I64(column) => Some(DynamicColumn::F64(Column {
|
||||
index: column.index,
|
||||
idx: column.idx,
|
||||
values: Arc::new(monotonic_map_column(column.values, MapI64ToF64)),
|
||||
})),
|
||||
DynamicColumn::U64(column) => Some(DynamicColumn::F64(Column {
|
||||
index: column.index,
|
||||
idx: column.idx,
|
||||
values: Arc::new(monotonic_map_column(column.values, MapU64ToF64)),
|
||||
})),
|
||||
DynamicColumn::F64(_) => Some(self),
|
||||
@@ -126,7 +91,7 @@ impl DynamicColumn {
|
||||
return None;
|
||||
}
|
||||
Some(DynamicColumn::I64(Column {
|
||||
index: column.index,
|
||||
idx: column.idx,
|
||||
values: Arc::new(monotonic_map_column(column.values, MapU64ToI64)),
|
||||
}))
|
||||
}
|
||||
@@ -141,7 +106,7 @@ impl DynamicColumn {
|
||||
return None;
|
||||
}
|
||||
Some(DynamicColumn::U64(Column {
|
||||
index: column.index,
|
||||
idx: column.idx,
|
||||
values: Arc::new(monotonic_map_column(column.values, MapI64ToU64)),
|
||||
}))
|
||||
}
|
||||
@@ -201,9 +166,9 @@ impl StrictlyMonotonicFn<i64, u64> for MapI64ToU64 {
|
||||
|
||||
macro_rules! static_dynamic_conversions {
|
||||
($typ:ty, $enum_name:ident) => {
|
||||
impl From<DynamicColumn> for Option<$typ> {
|
||||
fn from(dynamic_column: DynamicColumn) -> Option<$typ> {
|
||||
if let DynamicColumn::$enum_name(col) = dynamic_column {
|
||||
impl Into<Option<$typ>> for DynamicColumn {
|
||||
fn into(self) -> Option<$typ> {
|
||||
if let DynamicColumn::$enum_name(col) = self {
|
||||
Some(col)
|
||||
} else {
|
||||
None
|
||||
@@ -223,7 +188,7 @@ static_dynamic_conversions!(Column<bool>, Bool);
|
||||
static_dynamic_conversions!(Column<u64>, U64);
|
||||
static_dynamic_conversions!(Column<i64>, I64);
|
||||
static_dynamic_conversions!(Column<f64>, F64);
|
||||
static_dynamic_conversions!(Column<DateTime>, DateTime);
|
||||
static_dynamic_conversions!(Column<crate::DateTime>, DateTime);
|
||||
static_dynamic_conversions!(StrColumn, Str);
|
||||
static_dynamic_conversions!(BytesColumn, Bytes);
|
||||
static_dynamic_conversions!(Column<Ipv6Addr>, IpAddr);
|
||||
@@ -241,9 +206,10 @@ impl DynamicColumnHandle {
|
||||
self.open_internal(column_bytes)
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
pub fn file_slice(&self) -> &FileSlice {
|
||||
&self.file_slice
|
||||
// TODO rename load_async
|
||||
pub async fn open_async(&self) -> io::Result<DynamicColumn> {
|
||||
let column_bytes: OwnedBytes = self.file_slice.read_bytes_async().await?;
|
||||
self.open_internal(column_bytes)
|
||||
}
|
||||
|
||||
/// Returns the `u64` fast field reader reader associated with `fields` of types
|
||||
@@ -269,22 +235,24 @@ impl DynamicColumnHandle {
|
||||
|
||||
fn open_internal(&self, column_bytes: OwnedBytes) -> io::Result<DynamicColumn> {
|
||||
let dynamic_column: DynamicColumn = match self.column_type {
|
||||
ColumnType::Bytes => crate::column::open_column_bytes(column_bytes)?.into(),
|
||||
ColumnType::Str => crate::column::open_column_str(column_bytes)?.into(),
|
||||
ColumnType::Bytes => {
|
||||
crate::column::open_column_bytes::<BytesColumn>(column_bytes)?.into()
|
||||
}
|
||||
ColumnType::Str => crate::column::open_column_bytes::<StrColumn>(column_bytes)?.into(),
|
||||
ColumnType::I64 => crate::column::open_column_u64::<i64>(column_bytes)?.into(),
|
||||
ColumnType::U64 => crate::column::open_column_u64::<u64>(column_bytes)?.into(),
|
||||
ColumnType::F64 => crate::column::open_column_u64::<f64>(column_bytes)?.into(),
|
||||
ColumnType::Bool => crate::column::open_column_u64::<bool>(column_bytes)?.into(),
|
||||
ColumnType::IpAddr => crate::column::open_column_u128::<Ipv6Addr>(column_bytes)?.into(),
|
||||
ColumnType::DateTime => {
|
||||
crate::column::open_column_u64::<DateTime>(column_bytes)?.into()
|
||||
crate::column::open_column_u64::<crate::DateTime>(column_bytes)?.into()
|
||||
}
|
||||
};
|
||||
Ok(dynamic_column)
|
||||
}
|
||||
|
||||
pub fn num_bytes(&self) -> ByteCount {
|
||||
self.file_slice.len().into()
|
||||
pub fn num_bytes(&self) -> usize {
|
||||
self.file_slice.len()
|
||||
}
|
||||
|
||||
pub fn column_type(&self) -> ColumnType {
|
||||
|
||||
@@ -7,10 +7,8 @@ extern crate more_asserts;
|
||||
#[cfg(all(test, feature = "unstable"))]
|
||||
extern crate test;
|
||||
|
||||
use std::fmt::Display;
|
||||
use std::io;
|
||||
|
||||
mod block_accessor;
|
||||
mod column;
|
||||
mod column_index;
|
||||
pub mod column_values;
|
||||
@@ -21,15 +19,12 @@ mod iterable;
|
||||
pub(crate) mod utils;
|
||||
mod value;
|
||||
|
||||
pub use block_accessor::ColumnBlockAccessor;
|
||||
pub use column::{BytesColumn, Column, StrColumn};
|
||||
pub use column_index::ColumnIndex;
|
||||
pub use column_values::{
|
||||
ColumnValues, EmptyColumnValues, MonotonicallyMappableToU128, MonotonicallyMappableToU64,
|
||||
};
|
||||
pub use column_values::{ColumnValues, MonotonicallyMappableToU128, MonotonicallyMappableToU64};
|
||||
pub use columnar::{
|
||||
merge_columnar, ColumnType, ColumnarReader, ColumnarWriter, HasAssociatedColumnType,
|
||||
MergeRowOrder, ShuffleMergeOrder, StackMergeOrder,
|
||||
MergeRowOrder, StackMergeOrder,
|
||||
};
|
||||
use sstable::VoidSSTable;
|
||||
pub use value::{NumericalType, NumericalValue};
|
||||
@@ -37,18 +32,13 @@ pub use value::{NumericalType, NumericalValue};
|
||||
pub use self::dynamic_column::{DynamicColumn, DynamicColumnHandle};
|
||||
|
||||
pub type RowId = u32;
|
||||
pub type DocId = u32;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct RowAddr {
|
||||
pub segment_ord: u32,
|
||||
pub row_id: RowId,
|
||||
}
|
||||
|
||||
pub use sstable::Dictionary;
|
||||
pub type Streamer<'a> = sstable::Streamer<'a, VoidSSTable>;
|
||||
|
||||
pub use common::DateTime;
|
||||
#[derive(Clone, Copy, PartialOrd, PartialEq, Default, Debug)]
|
||||
pub struct DateTime {
|
||||
pub timestamp_micros: i64,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub struct InvalidData;
|
||||
@@ -76,17 +66,6 @@ pub enum Cardinality {
|
||||
Multivalued = 2,
|
||||
}
|
||||
|
||||
impl Display for Cardinality {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
let short_str = match self {
|
||||
Cardinality::Full => "full",
|
||||
Cardinality::Optional => "opt",
|
||||
Cardinality::Multivalued => "mult",
|
||||
};
|
||||
write!(f, "{short_str}")
|
||||
}
|
||||
}
|
||||
|
||||
impl Cardinality {
|
||||
pub fn is_optional(&self) -> bool {
|
||||
matches!(self, Cardinality::Optional)
|
||||
@@ -97,6 +76,7 @@ impl Cardinality {
|
||||
pub(crate) fn to_code(self) -> u8 {
|
||||
self as u8
|
||||
}
|
||||
|
||||
pub(crate) fn try_from_code(code: u8) -> Result<Cardinality, InvalidData> {
|
||||
match code {
|
||||
0 => Ok(Cardinality::Full),
|
||||
|
||||
@@ -1,17 +1,10 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Debug;
|
||||
use std::net::Ipv6Addr;
|
||||
|
||||
use common::DateTime;
|
||||
use proptest::prelude::*;
|
||||
|
||||
use crate::column_values::MonotonicallyMappableToU128;
|
||||
use crate::columnar::ColumnType;
|
||||
use crate::dynamic_column::{DynamicColumn, DynamicColumnHandle};
|
||||
use crate::value::{Coerce, NumericalValue};
|
||||
use crate::{
|
||||
BytesColumn, Cardinality, Column, ColumnarReader, ColumnarWriter, RowId, StackMergeOrder,
|
||||
};
|
||||
use crate::value::NumericalValue;
|
||||
use crate::{Cardinality, ColumnarReader, ColumnarWriter};
|
||||
|
||||
#[test]
|
||||
fn test_dataframe_writer_str() {
|
||||
@@ -19,12 +12,12 @@ fn test_dataframe_writer_str() {
|
||||
dataframe_writer.record_str(1u32, "my_string", "hello");
|
||||
dataframe_writer.record_str(3u32, "my_string", "helloeee");
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(5, None, &mut buffer).unwrap();
|
||||
dataframe_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("my_string").unwrap();
|
||||
assert_eq!(cols.len(), 1);
|
||||
assert_eq!(cols[0].num_bytes(), 89);
|
||||
assert_eq!(cols[0].num_bytes(), 158);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -33,12 +26,12 @@ fn test_dataframe_writer_bytes() {
|
||||
dataframe_writer.record_bytes(1u32, "my_string", b"hello");
|
||||
dataframe_writer.record_bytes(3u32, "my_string", b"helloeee");
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(5, None, &mut buffer).unwrap();
|
||||
dataframe_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("my_string").unwrap();
|
||||
assert_eq!(cols.len(), 1);
|
||||
assert_eq!(cols[0].num_bytes(), 89);
|
||||
assert_eq!(cols[0].num_bytes(), 158);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -47,7 +40,7 @@ fn test_dataframe_writer_bool() {
|
||||
dataframe_writer.record_bool(1u32, "bool.value", false);
|
||||
dataframe_writer.record_bool(3u32, "bool.value", true);
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(5, None, &mut buffer).unwrap();
|
||||
dataframe_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("bool.value").unwrap();
|
||||
@@ -70,7 +63,7 @@ fn test_dataframe_writer_u64_multivalued() {
|
||||
dataframe_writer.record_numerical(6u32, "divisor", 2u64);
|
||||
dataframe_writer.record_numerical(6u32, "divisor", 3u64);
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(7, None, &mut buffer).unwrap();
|
||||
dataframe_writer.serialize(7, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("divisor").unwrap();
|
||||
@@ -82,7 +75,7 @@ fn test_dataframe_writer_u64_multivalued() {
|
||||
divisor_col.get_cardinality(),
|
||||
crate::Cardinality::Multivalued
|
||||
);
|
||||
assert_eq!(divisor_col.num_docs(), 7);
|
||||
assert_eq!(divisor_col.num_rows(), 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -91,7 +84,7 @@ fn test_dataframe_writer_ip_addr() {
|
||||
dataframe_writer.record_ip_addr(1, "ip_addr", Ipv6Addr::from_u128(1001));
|
||||
dataframe_writer.record_ip_addr(3, "ip_addr", Ipv6Addr::from_u128(1050));
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(5, None, &mut buffer).unwrap();
|
||||
dataframe_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("ip_addr").unwrap();
|
||||
@@ -120,7 +113,7 @@ fn test_dataframe_writer_numerical() {
|
||||
dataframe_writer.record_numerical(2u32, "srical.value", NumericalValue::U64(13u64));
|
||||
dataframe_writer.record_numerical(4u32, "srical.value", NumericalValue::U64(15u64));
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(6, None, &mut buffer).unwrap();
|
||||
dataframe_writer.serialize(6, &mut buffer).unwrap();
|
||||
let columnar = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("srical.value").unwrap();
|
||||
@@ -133,7 +126,7 @@ fn test_dataframe_writer_numerical() {
|
||||
assert_eq!(cols[0].num_bytes(), 33);
|
||||
let column = cols[0].open().unwrap();
|
||||
let DynamicColumn::I64(column_i64) = column else { panic!(); };
|
||||
assert_eq!(column_i64.index.get_cardinality(), Cardinality::Optional);
|
||||
assert_eq!(column_i64.idx.get_cardinality(), Cardinality::Optional);
|
||||
assert_eq!(column_i64.first(0), None);
|
||||
assert_eq!(column_i64.first(1), Some(12i64));
|
||||
assert_eq!(column_i64.first(2), Some(13i64));
|
||||
@@ -151,7 +144,7 @@ fn test_dictionary_encoded_str() {
|
||||
columnar_writer.record_str(3, "my.column", "c");
|
||||
columnar_writer.record_str(3, "my.column2", "different_column!");
|
||||
columnar_writer.record_str(4, "my.column", "b");
|
||||
columnar_writer.serialize(5, None, &mut buffer).unwrap();
|
||||
columnar_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar_reader = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar_reader.num_columns(), 2);
|
||||
let col_handles = columnar_reader.read_columns("my.column").unwrap();
|
||||
@@ -183,7 +176,7 @@ fn test_dictionary_encoded_bytes() {
|
||||
columnar_writer.record_bytes(3, "my.column", b"c");
|
||||
columnar_writer.record_bytes(3, "my.column2", b"different_column!");
|
||||
columnar_writer.record_bytes(4, "my.column", b"b");
|
||||
columnar_writer.serialize(5, None, &mut buffer).unwrap();
|
||||
columnar_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar_reader = ColumnarReader::open(buffer).unwrap();
|
||||
assert_eq!(columnar_reader.num_columns(), 2);
|
||||
let col_handles = columnar_reader.read_columns("my.column").unwrap();
|
||||
@@ -217,422 +210,3 @@ fn test_dictionary_encoded_bytes() {
|
||||
.unwrap();
|
||||
assert_eq!(term_buffer, b"b");
|
||||
}
|
||||
|
||||
fn num_strategy() -> impl Strategy<Value = NumericalValue> {
|
||||
prop_oneof![
|
||||
Just(NumericalValue::U64(0u64)),
|
||||
Just(NumericalValue::U64(u64::MAX)),
|
||||
Just(NumericalValue::I64(0i64)),
|
||||
Just(NumericalValue::I64(i64::MIN)),
|
||||
Just(NumericalValue::I64(i64::MAX)),
|
||||
Just(NumericalValue::F64(1.2f64)),
|
||||
]
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum ColumnValue {
|
||||
Str(&'static str),
|
||||
Bytes(&'static [u8]),
|
||||
Numerical(NumericalValue),
|
||||
IpAddr(Ipv6Addr),
|
||||
Bool(bool),
|
||||
DateTime(DateTime),
|
||||
}
|
||||
|
||||
impl ColumnValue {
|
||||
pub(crate) fn column_type_category(&self) -> ColumnTypeCategory {
|
||||
match self {
|
||||
ColumnValue::Str(_) => ColumnTypeCategory::Str,
|
||||
ColumnValue::Bytes(_) => ColumnTypeCategory::Bytes,
|
||||
ColumnValue::Numerical(numerical_val) => ColumnTypeCategory::Numerical,
|
||||
ColumnValue::IpAddr(_) => ColumnTypeCategory::IpAddr,
|
||||
ColumnValue::Bool(_) => ColumnTypeCategory::Bool,
|
||||
ColumnValue::DateTime(_) => ColumnTypeCategory::DateTime,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn column_name_strategy() -> impl Strategy<Value = &'static str> {
|
||||
prop_oneof![Just("c1"), Just("c2")]
|
||||
}
|
||||
|
||||
fn string_strategy() -> impl Strategy<Value = &'static str> {
|
||||
prop_oneof![Just("a"), Just("b")]
|
||||
}
|
||||
|
||||
fn bytes_strategy() -> impl Strategy<Value = &'static [u8]> {
|
||||
prop_oneof![Just(&[0u8][..]), Just(&[1u8][..])]
|
||||
}
|
||||
|
||||
// A random column value
|
||||
fn column_value_strategy() -> impl Strategy<Value = ColumnValue> {
|
||||
prop_oneof![
|
||||
string_strategy().prop_map(|s| ColumnValue::Str(s)),
|
||||
bytes_strategy().prop_map(|b| ColumnValue::Bytes(b)),
|
||||
num_strategy().prop_map(|n| ColumnValue::Numerical(n)),
|
||||
(1u16..3u16).prop_map(|ip_addr_byte| ColumnValue::IpAddr(Ipv6Addr::new(
|
||||
127,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
ip_addr_byte
|
||||
))),
|
||||
any::<bool>().prop_map(|b| ColumnValue::Bool(b)),
|
||||
(0_679_723_993i64..1_679_723_995i64)
|
||||
.prop_map(|val| { ColumnValue::DateTime(DateTime::from_timestamp_secs(val)) })
|
||||
]
|
||||
}
|
||||
|
||||
// A document contains up to 4 values.
|
||||
fn doc_strategy() -> impl Strategy<Value = Vec<(&'static str, ColumnValue)>> {
|
||||
proptest::collection::vec((column_name_strategy(), column_value_strategy()), 0..4)
|
||||
}
|
||||
|
||||
// A columnar contains up to 2 docs.
|
||||
fn columnar_docs_strategy() -> impl Strategy<Value = Vec<Vec<(&'static str, ColumnValue)>>> {
|
||||
proptest::collection::vec(doc_strategy(), 0..=2)
|
||||
}
|
||||
|
||||
fn columnar_docs_and_mapping_strategy(
|
||||
) -> impl Strategy<Value = (Vec<Vec<(&'static str, ColumnValue)>>, Vec<RowId>)> {
|
||||
columnar_docs_strategy().prop_flat_map(|docs| {
|
||||
permutation_strategy(docs.len()).prop_map(move |permutation| (docs.clone(), permutation))
|
||||
})
|
||||
}
|
||||
|
||||
fn permutation_strategy(n: usize) -> impl Strategy<Value = Vec<RowId>> {
|
||||
Just((0u32..n as RowId).collect()).prop_shuffle()
|
||||
}
|
||||
|
||||
fn build_columnar_with_mapping(
|
||||
docs: &[Vec<(&'static str, ColumnValue)>],
|
||||
old_to_new_row_ids_opt: Option<&[RowId]>,
|
||||
) -> ColumnarReader {
|
||||
let num_docs = docs.len() as u32;
|
||||
let mut buffer = Vec::new();
|
||||
let mut columnar_writer = ColumnarWriter::default();
|
||||
for (doc_id, vals) in docs.iter().enumerate() {
|
||||
for (column_name, col_val) in vals {
|
||||
match *col_val {
|
||||
ColumnValue::Str(str_val) => {
|
||||
columnar_writer.record_str(doc_id as u32, column_name, str_val);
|
||||
}
|
||||
ColumnValue::Bytes(bytes) => {
|
||||
columnar_writer.record_bytes(doc_id as u32, column_name, bytes)
|
||||
}
|
||||
ColumnValue::Numerical(num) => {
|
||||
columnar_writer.record_numerical(doc_id as u32, column_name, num);
|
||||
}
|
||||
ColumnValue::IpAddr(ip_addr) => {
|
||||
columnar_writer.record_ip_addr(doc_id as u32, column_name, ip_addr);
|
||||
}
|
||||
ColumnValue::Bool(bool_val) => {
|
||||
columnar_writer.record_bool(doc_id as u32, column_name, bool_val);
|
||||
}
|
||||
ColumnValue::DateTime(date_time) => {
|
||||
columnar_writer.record_datetime(doc_id as u32, column_name, date_time);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
columnar_writer
|
||||
.serialize(num_docs, old_to_new_row_ids_opt, &mut buffer)
|
||||
.unwrap();
|
||||
let columnar_reader = ColumnarReader::open(buffer).unwrap();
|
||||
columnar_reader
|
||||
}
|
||||
|
||||
fn build_columnar(docs: &[Vec<(&'static str, ColumnValue)>]) -> ColumnarReader {
|
||||
build_columnar_with_mapping(docs, None)
|
||||
}
|
||||
|
||||
fn assert_columnar_eq(left: &ColumnarReader, right: &ColumnarReader) {
|
||||
assert_eq!(left.num_rows(), right.num_rows());
|
||||
let left_columns = left.list_columns().unwrap();
|
||||
let right_columns = right.list_columns().unwrap();
|
||||
assert_eq!(left_columns.len(), right_columns.len());
|
||||
for i in 0..left_columns.len() {
|
||||
assert_eq!(left_columns[i].0, right_columns[i].0);
|
||||
let left_column = left_columns[i].1.open().unwrap();
|
||||
let right_column = right_columns[i].1.open().unwrap();
|
||||
assert_dyn_column_eq(&left_column, &right_column);
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_column_eq<T: PartialEq + Copy>(left: &Column<T>, right: &Column<T>) {}
|
||||
|
||||
fn assert_bytes_column_eq(left: &BytesColumn, right: &BytesColumn) {}
|
||||
|
||||
fn assert_dyn_column_eq(left_dyn_column: &DynamicColumn, right_dyn_column: &DynamicColumn) {
|
||||
assert_eq!(
|
||||
&left_dyn_column.column_type(),
|
||||
&right_dyn_column.column_type()
|
||||
);
|
||||
assert_eq!(
|
||||
&left_dyn_column.get_cardinality(),
|
||||
&right_dyn_column.get_cardinality()
|
||||
);
|
||||
match &(left_dyn_column, right_dyn_column) {
|
||||
(DynamicColumn::Bool(left_col), DynamicColumn::Bool(right_col)) => {
|
||||
assert_column_eq(left_col, right_col);
|
||||
}
|
||||
(DynamicColumn::I64(left_col), DynamicColumn::I64(right_col)) => {
|
||||
assert_column_eq(left_col, right_col);
|
||||
}
|
||||
(DynamicColumn::U64(left_col), DynamicColumn::U64(right_col)) => {
|
||||
assert_column_eq(left_col, right_col);
|
||||
}
|
||||
(DynamicColumn::F64(left_col), DynamicColumn::F64(right_col)) => {
|
||||
assert_column_eq(left_col, right_col);
|
||||
}
|
||||
(DynamicColumn::DateTime(left_col), DynamicColumn::DateTime(right_col)) => {
|
||||
assert_column_eq(left_col, right_col);
|
||||
}
|
||||
(DynamicColumn::IpAddr(left_col), DynamicColumn::IpAddr(right_col)) => {
|
||||
assert_column_eq(left_col, right_col);
|
||||
}
|
||||
(DynamicColumn::Bytes(left_col), DynamicColumn::Bytes(right_col)) => {
|
||||
assert_bytes_column_eq(left_col, right_col);
|
||||
}
|
||||
(DynamicColumn::Str(left_col), DynamicColumn::Str(right_col)) => {
|
||||
assert_bytes_column_eq(left_col, right_col);
|
||||
}
|
||||
_ => {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
trait AssertEqualToColumnValue {
|
||||
fn assert_equal_to_column_value(&self, column_value: &ColumnValue);
|
||||
}
|
||||
|
||||
use crate::columnar::ColumnTypeCategory;
|
||||
|
||||
impl AssertEqualToColumnValue for bool {
|
||||
fn assert_equal_to_column_value(&self, column_value: &ColumnValue) {
|
||||
let ColumnValue::Bool(val) = column_value else { panic!() };
|
||||
assert_eq!(self, val);
|
||||
}
|
||||
}
|
||||
|
||||
impl AssertEqualToColumnValue for Ipv6Addr {
|
||||
fn assert_equal_to_column_value(&self, column_value: &ColumnValue) {
|
||||
let ColumnValue::IpAddr(val) = column_value else { panic!() };
|
||||
assert_eq!(self, val);
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Coerce + PartialEq + Debug + Into<NumericalValue>> AssertEqualToColumnValue for T {
|
||||
fn assert_equal_to_column_value(&self, column_value: &ColumnValue) {
|
||||
let ColumnValue::Numerical(num) = column_value else { panic!() };
|
||||
assert_eq!(self, &T::coerce(*num));
|
||||
}
|
||||
}
|
||||
|
||||
impl AssertEqualToColumnValue for DateTime {
|
||||
fn assert_equal_to_column_value(&self, column_value: &ColumnValue) {
|
||||
let ColumnValue::DateTime(dt) = column_value else { panic!() };
|
||||
assert_eq!(self, dt);
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_column_values<
|
||||
T: AssertEqualToColumnValue + PartialEq + Copy + PartialOrd + Debug + Send + Sync + 'static,
|
||||
>(
|
||||
col: &Column<T>,
|
||||
expected: &HashMap<u32, Vec<&ColumnValue>>,
|
||||
) {
|
||||
let mut num_non_empty_rows = 0;
|
||||
for doc in 0..col.num_docs() {
|
||||
let doc_vals: Vec<T> = col.values_for_doc(doc).collect();
|
||||
if doc_vals.is_empty() {
|
||||
continue;
|
||||
}
|
||||
num_non_empty_rows += 1;
|
||||
let expected_vals = expected.get(&doc).unwrap();
|
||||
assert_eq!(doc_vals.len(), expected_vals.len());
|
||||
for (val, &expected) in doc_vals.iter().zip(expected_vals.iter()) {
|
||||
val.assert_equal_to_column_value(expected)
|
||||
}
|
||||
}
|
||||
assert_eq!(num_non_empty_rows, expected.len());
|
||||
}
|
||||
|
||||
fn assert_bytes_column_values(
|
||||
col: &BytesColumn,
|
||||
expected: &HashMap<u32, Vec<&ColumnValue>>,
|
||||
is_str: bool,
|
||||
) {
|
||||
let mut num_non_empty_rows = 0;
|
||||
let mut buffer = Vec::new();
|
||||
for doc in 0..col.term_ord_column.num_docs() {
|
||||
let doc_vals: Vec<u64> = col.term_ords(doc).collect();
|
||||
if doc_vals.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let expected_vals = expected.get(&doc).unwrap();
|
||||
assert_eq!(doc_vals.len(), expected_vals.len());
|
||||
for (&expected_col_val, &ord) in expected_vals.iter().zip(&doc_vals) {
|
||||
col.ord_to_bytes(ord, &mut buffer).unwrap();
|
||||
match expected_col_val {
|
||||
ColumnValue::Str(str_val) => {
|
||||
assert!(is_str);
|
||||
assert_eq!(str_val.as_bytes(), &buffer);
|
||||
}
|
||||
ColumnValue::Bytes(bytes_val) => {
|
||||
assert!(!is_str);
|
||||
assert_eq!(bytes_val, &buffer);
|
||||
}
|
||||
_ => {
|
||||
panic!();
|
||||
}
|
||||
}
|
||||
}
|
||||
num_non_empty_rows += 1;
|
||||
}
|
||||
assert_eq!(num_non_empty_rows, expected.len());
|
||||
}
|
||||
|
||||
proptest! {
|
||||
/// This proptest attempts to create a tiny columnar based of up to 3 rows, and checks that the resulting
|
||||
/// columnar matches the row data.
|
||||
#[test]
|
||||
fn test_single_columnar_builder_proptest(docs in columnar_docs_strategy()) {
|
||||
let columnar = build_columnar(&docs[..]);
|
||||
assert_eq!(columnar.num_rows() as usize, docs.len());
|
||||
let mut expected_columns: HashMap<(&str, ColumnTypeCategory), HashMap<u32, Vec<&ColumnValue>> > = Default::default();
|
||||
for (doc_id, doc_vals) in docs.iter().enumerate() {
|
||||
for (col_name, col_val) in doc_vals {
|
||||
expected_columns
|
||||
.entry((col_name, col_val.column_type_category()))
|
||||
.or_default()
|
||||
.entry(doc_id as u32)
|
||||
.or_default()
|
||||
.push(col_val);
|
||||
}
|
||||
}
|
||||
let column_list = columnar.list_columns().unwrap();
|
||||
assert_eq!(expected_columns.len(), column_list.len());
|
||||
for (column_name, column) in column_list {
|
||||
let dynamic_column = column.open().unwrap();
|
||||
let col_category: ColumnTypeCategory = dynamic_column.column_type().into();
|
||||
let expected_col_values: &HashMap<u32, Vec<&ColumnValue>> = expected_columns.get(&(column_name.as_str(), col_category)).unwrap();
|
||||
match &dynamic_column {
|
||||
DynamicColumn::Bool(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::I64(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::U64(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::F64(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::IpAddr(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::DateTime(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::Bytes(col) =>
|
||||
assert_bytes_column_values(col, expected_col_values, false),
|
||||
DynamicColumn::Str(col) =>
|
||||
assert_bytes_column_values(col, expected_col_values, true),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Same as `test_single_columnar_builder_proptest` but with a shuffling mapping.
|
||||
#[test]
|
||||
fn test_single_columnar_builder_with_shuffle_proptest((docs, mapping) in columnar_docs_and_mapping_strategy()) {
|
||||
let columnar = build_columnar_with_mapping(&docs[..], Some(&mapping));
|
||||
assert_eq!(columnar.num_rows() as usize, docs.len());
|
||||
let mut expected_columns: HashMap<(&str, ColumnTypeCategory), HashMap<u32, Vec<&ColumnValue>> > = Default::default();
|
||||
for (doc_id, doc_vals) in docs.iter().enumerate() {
|
||||
for (col_name, col_val) in doc_vals {
|
||||
expected_columns
|
||||
.entry((col_name, col_val.column_type_category()))
|
||||
.or_default()
|
||||
.entry(mapping[doc_id])
|
||||
.or_default()
|
||||
.push(col_val);
|
||||
}
|
||||
}
|
||||
let column_list = columnar.list_columns().unwrap();
|
||||
assert_eq!(expected_columns.len(), column_list.len());
|
||||
for (column_name, column) in column_list {
|
||||
let dynamic_column = column.open().unwrap();
|
||||
let col_category: ColumnTypeCategory = dynamic_column.column_type().into();
|
||||
let expected_col_values: &HashMap<u32, Vec<&ColumnValue>> = expected_columns.get(&(column_name.as_str(), col_category)).unwrap();
|
||||
for doc_id in 0..columnar.num_rows() {
|
||||
match &dynamic_column {
|
||||
DynamicColumn::Bool(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::I64(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::U64(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::F64(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::IpAddr(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::DateTime(col) =>
|
||||
assert_column_values(col, expected_col_values),
|
||||
DynamicColumn::Bytes(col) =>
|
||||
assert_bytes_column_values(col, expected_col_values, false),
|
||||
DynamicColumn::Str(col) =>
|
||||
assert_bytes_column_values(col, expected_col_values, true),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// This tests create 2 or 3 random small columnar and attempts to merge them.
|
||||
/// It compares the resulting merged dataframe with what would have been obtained by building the
|
||||
/// dataframe from the concatenated rows to begin with.
|
||||
#[test]
|
||||
fn test_columnar_merge_proptest(columnar_docs in proptest::collection::vec(columnar_docs_strategy(), 2..=3)) {
|
||||
let columnar_readers: Vec<ColumnarReader> = columnar_docs.iter()
|
||||
.map(|docs| build_columnar(&docs[..]))
|
||||
.collect::<Vec<_>>();
|
||||
let columnar_readers_arr: Vec<&ColumnarReader> = columnar_readers.iter().collect();
|
||||
let mut output: Vec<u8> = Vec::new();
|
||||
let stack_merge_order = StackMergeOrder::stack(&columnar_readers_arr[..]);
|
||||
crate::merge_columnar(&columnar_readers_arr[..], &[], crate::MergeRowOrder::Stack(stack_merge_order), &mut output).unwrap();
|
||||
let merged_columnar = ColumnarReader::open(output).unwrap();
|
||||
let concat_rows: Vec<Vec<(&'static str, ColumnValue)>> = columnar_docs.iter().cloned().flatten().collect();
|
||||
let expected_merged_columnar = build_columnar(&concat_rows[..]);
|
||||
assert_columnar_eq(&merged_columnar, &expected_merged_columnar);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_columnar_failing_test() {
|
||||
let columnar_docs: Vec<Vec<Vec<(&str, ColumnValue)>>> =
|
||||
vec![vec![], vec![vec![("c1", ColumnValue::Str("a"))]]];
|
||||
let columnar_readers: Vec<ColumnarReader> = columnar_docs
|
||||
.iter()
|
||||
.map(|docs| build_columnar(&docs[..]))
|
||||
.collect::<Vec<_>>();
|
||||
let columnar_readers_arr: Vec<&ColumnarReader> = columnar_readers.iter().collect();
|
||||
let mut output: Vec<u8> = Vec::new();
|
||||
let stack_merge_order = StackMergeOrder::stack(&columnar_readers_arr[..]);
|
||||
crate::merge_columnar(
|
||||
&columnar_readers_arr[..],
|
||||
&[],
|
||||
crate::MergeRowOrder::Stack(stack_merge_order),
|
||||
&mut output,
|
||||
)
|
||||
.unwrap();
|
||||
let merged_columnar = ColumnarReader::open(output).unwrap();
|
||||
let concat_rows: Vec<Vec<(&'static str, ColumnValue)>> =
|
||||
columnar_docs.iter().cloned().flatten().collect();
|
||||
let expected_merged_columnar = build_columnar(&concat_rows[..]);
|
||||
assert_columnar_eq(&merged_columnar, &expected_merged_columnar);
|
||||
}
|
||||
|
||||
// TODO add non trivial remap and merge
|
||||
// TODO test required
|
||||
// TODO add support for empty columnar.
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
use common::DateTime;
|
||||
|
||||
use crate::InvalidData;
|
||||
|
||||
#[derive(Copy, Clone, PartialEq, Debug)]
|
||||
@@ -106,10 +104,10 @@ impl Coerce for f64 {
|
||||
}
|
||||
}
|
||||
|
||||
impl Coerce for DateTime {
|
||||
impl Coerce for crate::DateTime {
|
||||
fn coerce(value: NumericalValue) -> Self {
|
||||
let timestamp_micros = i64::coerce(value);
|
||||
DateTime::from_timestamp_micros(timestamp_micros)
|
||||
crate::DateTime { timestamp_micros }
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,8 +16,6 @@ repository = "https://github.com/quickwit-oss/tantivy"
|
||||
byteorder = "1.4.3"
|
||||
ownedbytes = { version= "0.5", path="../ownedbytes" }
|
||||
async-trait = "0.1"
|
||||
time = { version = "0.3.10", features = ["serde-well-known"] }
|
||||
serde = { version = "1.0.136", features = ["derive"] }
|
||||
|
||||
[dev-dependencies]
|
||||
proptest = "1.0.0"
|
||||
|
||||
@@ -4,8 +4,6 @@ use std::{fmt, io, u64};
|
||||
|
||||
use ownedbytes::OwnedBytes;
|
||||
|
||||
use crate::ByteCount;
|
||||
|
||||
#[derive(Clone, Copy, Eq, PartialEq)]
|
||||
pub struct TinySet(u64);
|
||||
|
||||
@@ -388,8 +386,8 @@ impl ReadOnlyBitSet {
|
||||
}
|
||||
|
||||
/// Number of bytes used in the bitset representation.
|
||||
pub fn num_bytes(&self) -> ByteCount {
|
||||
self.data.len().into()
|
||||
pub fn num_bytes(&self) -> usize {
|
||||
self.data.len()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,114 +0,0 @@
|
||||
use std::iter::Sum;
|
||||
use std::ops::{Add, AddAssign};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Indicates space usage in bytes
|
||||
#[derive(Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
|
||||
pub struct ByteCount(u64);
|
||||
|
||||
impl std::fmt::Debug for ByteCount {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(&self.human_readable())
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ByteCount {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(&self.human_readable())
|
||||
}
|
||||
}
|
||||
|
||||
const SUFFIX_AND_THRESHOLD: [(&str, u64); 5] = [
|
||||
("KB", 1_000),
|
||||
("MB", 1_000_000),
|
||||
("GB", 1_000_000_000),
|
||||
("TB", 1_000_000_000_000),
|
||||
("PB", 1_000_000_000_000_000),
|
||||
];
|
||||
|
||||
impl ByteCount {
|
||||
#[inline]
|
||||
pub fn get_bytes(&self) -> u64 {
|
||||
self.0
|
||||
}
|
||||
|
||||
pub fn human_readable(&self) -> String {
|
||||
for (suffix, threshold) in SUFFIX_AND_THRESHOLD.iter().rev() {
|
||||
if self.get_bytes() >= *threshold {
|
||||
let unit_num = self.get_bytes() as f64 / *threshold as f64;
|
||||
return format!("{:.2} {}", unit_num, suffix);
|
||||
}
|
||||
}
|
||||
format!("{:.2} B", self.get_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<u64> for ByteCount {
|
||||
fn from(value: u64) -> Self {
|
||||
ByteCount(value)
|
||||
}
|
||||
}
|
||||
impl From<usize> for ByteCount {
|
||||
fn from(value: usize) -> Self {
|
||||
ByteCount(value as u64)
|
||||
}
|
||||
}
|
||||
|
||||
impl Sum for ByteCount {
|
||||
#[inline]
|
||||
fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
|
||||
iter.fold(ByteCount::default(), |acc, x| acc + x)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<u64> for ByteCount {
|
||||
#[inline]
|
||||
fn eq(&self, other: &u64) -> bool {
|
||||
self.get_bytes() == *other
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialOrd<u64> for ByteCount {
|
||||
#[inline]
|
||||
fn partial_cmp(&self, other: &u64) -> Option<std::cmp::Ordering> {
|
||||
self.get_bytes().partial_cmp(other)
|
||||
}
|
||||
}
|
||||
|
||||
impl Add for ByteCount {
|
||||
type Output = Self;
|
||||
|
||||
#[inline]
|
||||
fn add(self, other: Self) -> Self {
|
||||
Self(self.get_bytes() + other.get_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
impl AddAssign for ByteCount {
|
||||
#[inline]
|
||||
fn add_assign(&mut self, other: Self) {
|
||||
*self = Self(self.get_bytes() + other.get_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use crate::ByteCount;
|
||||
|
||||
#[test]
|
||||
fn test_bytes() {
|
||||
assert_eq!(ByteCount::from(0u64).human_readable(), "0 B");
|
||||
assert_eq!(ByteCount::from(300u64).human_readable(), "300 B");
|
||||
assert_eq!(ByteCount::from(1_000_000u64).human_readable(), "1.00 MB");
|
||||
assert_eq!(ByteCount::from(1_500_000u64).human_readable(), "1.50 MB");
|
||||
assert_eq!(
|
||||
ByteCount::from(1_500_000_000u64).human_readable(),
|
||||
"1.50 GB"
|
||||
);
|
||||
assert_eq!(
|
||||
ByteCount::from(3_213_000_000_000u64).human_readable(),
|
||||
"3.21 TB"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,136 +0,0 @@
|
||||
use std::fmt;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use time::format_description::well_known::Rfc3339;
|
||||
use time::{OffsetDateTime, PrimitiveDateTime, UtcOffset};
|
||||
|
||||
/// DateTime Precision
|
||||
#[derive(
|
||||
Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default,
|
||||
)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum DatePrecision {
|
||||
/// Seconds precision
|
||||
#[default]
|
||||
Seconds,
|
||||
/// Milli-seconds precision.
|
||||
Milliseconds,
|
||||
/// Micro-seconds precision.
|
||||
Microseconds,
|
||||
}
|
||||
|
||||
/// A date/time value with microsecond precision.
|
||||
///
|
||||
/// This timestamp does not carry any explicit time zone information.
|
||||
/// Users are responsible for applying the provided conversion
|
||||
/// functions consistently. Internally the time zone is assumed
|
||||
/// to be UTC, which is also used implicitly for JSON serialization.
|
||||
///
|
||||
/// All constructors and conversions are provided as explicit
|
||||
/// functions and not by implementing any `From`/`Into` traits
|
||||
/// to prevent unintended usage.
|
||||
#[derive(Clone, Default, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct DateTime {
|
||||
// Timestamp in microseconds.
|
||||
pub(crate) timestamp_micros: i64,
|
||||
}
|
||||
|
||||
impl DateTime {
|
||||
/// Create new from UNIX timestamp in seconds
|
||||
pub const fn from_timestamp_secs(seconds: i64) -> Self {
|
||||
Self {
|
||||
timestamp_micros: seconds * 1_000_000,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create new from UNIX timestamp in milliseconds
|
||||
pub const fn from_timestamp_millis(milliseconds: i64) -> Self {
|
||||
Self {
|
||||
timestamp_micros: milliseconds * 1_000,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create new from UNIX timestamp in microseconds.
|
||||
pub const fn from_timestamp_micros(microseconds: i64) -> Self {
|
||||
Self {
|
||||
timestamp_micros: microseconds,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create new from `OffsetDateTime`
|
||||
///
|
||||
/// The given date/time is converted to UTC and the actual
|
||||
/// time zone is discarded.
|
||||
pub const fn from_utc(dt: OffsetDateTime) -> Self {
|
||||
let timestamp_micros = dt.unix_timestamp() * 1_000_000 + dt.microsecond() as i64;
|
||||
Self { timestamp_micros }
|
||||
}
|
||||
|
||||
/// Create new from `PrimitiveDateTime`
|
||||
///
|
||||
/// Implicitly assumes that the given date/time is in UTC!
|
||||
/// Otherwise the original value must only be reobtained with
|
||||
/// [`Self::into_primitive()`].
|
||||
pub fn from_primitive(dt: PrimitiveDateTime) -> Self {
|
||||
Self::from_utc(dt.assume_utc())
|
||||
}
|
||||
|
||||
/// Convert to UNIX timestamp in seconds.
|
||||
pub const fn into_timestamp_secs(self) -> i64 {
|
||||
self.timestamp_micros / 1_000_000
|
||||
}
|
||||
|
||||
/// Convert to UNIX timestamp in milliseconds.
|
||||
pub const fn into_timestamp_millis(self) -> i64 {
|
||||
self.timestamp_micros / 1_000
|
||||
}
|
||||
|
||||
/// Convert to UNIX timestamp in microseconds.
|
||||
pub const fn into_timestamp_micros(self) -> i64 {
|
||||
self.timestamp_micros
|
||||
}
|
||||
|
||||
/// Convert to UTC `OffsetDateTime`
|
||||
pub fn into_utc(self) -> OffsetDateTime {
|
||||
let timestamp_nanos = self.timestamp_micros as i128 * 1000;
|
||||
let utc_datetime = OffsetDateTime::from_unix_timestamp_nanos(timestamp_nanos)
|
||||
.expect("valid UNIX timestamp");
|
||||
debug_assert_eq!(UtcOffset::UTC, utc_datetime.offset());
|
||||
utc_datetime
|
||||
}
|
||||
|
||||
/// Convert to `OffsetDateTime` with the given time zone
|
||||
pub fn into_offset(self, offset: UtcOffset) -> OffsetDateTime {
|
||||
self.into_utc().to_offset(offset)
|
||||
}
|
||||
|
||||
/// Convert to `PrimitiveDateTime` without any time zone
|
||||
///
|
||||
/// The value should have been constructed with [`Self::from_primitive()`].
|
||||
/// Otherwise the time zone is implicitly assumed to be UTC.
|
||||
pub fn into_primitive(self) -> PrimitiveDateTime {
|
||||
let utc_datetime = self.into_utc();
|
||||
// Discard the UTC time zone offset
|
||||
debug_assert_eq!(UtcOffset::UTC, utc_datetime.offset());
|
||||
PrimitiveDateTime::new(utc_datetime.date(), utc_datetime.time())
|
||||
}
|
||||
|
||||
/// Truncates the microseconds value to the corresponding precision.
|
||||
pub fn truncate(self, precision: DatePrecision) -> Self {
|
||||
let truncated_timestamp_micros = match precision {
|
||||
DatePrecision::Seconds => (self.timestamp_micros / 1_000_000) * 1_000_000,
|
||||
DatePrecision::Milliseconds => (self.timestamp_micros / 1_000) * 1_000,
|
||||
DatePrecision::Microseconds => self.timestamp_micros,
|
||||
};
|
||||
Self {
|
||||
timestamp_micros: truncated_timestamp_micros,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for DateTime {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let utc_rfc3339 = self.into_utc().format(&Rfc3339).map_err(|_| fmt::Error)?;
|
||||
f.write_str(&utc_rfc3339)
|
||||
}
|
||||
}
|
||||
@@ -1,63 +0,0 @@
|
||||
use std::io::{self, Read, Write};
|
||||
|
||||
use crate::BinarySerializable;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[repr(u32)]
|
||||
pub enum DictionaryKind {
|
||||
Fst = 1,
|
||||
SSTable = 2,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DictionaryFooter {
|
||||
pub kind: DictionaryKind,
|
||||
pub version: u32,
|
||||
}
|
||||
|
||||
impl DictionaryFooter {
|
||||
pub fn verify_equal(&self, other: &DictionaryFooter) -> io::Result<()> {
|
||||
if self.kind != other.kind {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::Other,
|
||||
format!(
|
||||
"Invalid dictionary type, expected {:?}, found {:?}",
|
||||
self.kind, other.kind
|
||||
),
|
||||
));
|
||||
}
|
||||
if self.version != other.version {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::Other,
|
||||
format!(
|
||||
"Unsuported dictionary version, expected {}, found {}",
|
||||
self.version, other.version
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl BinarySerializable for DictionaryFooter {
|
||||
fn serialize<W: Write + ?Sized>(&self, writer: &mut W) -> io::Result<()> {
|
||||
self.version.serialize(writer)?;
|
||||
(self.kind as u32).serialize(writer)
|
||||
}
|
||||
fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
|
||||
let version = u32::deserialize(reader)?;
|
||||
let kind = u32::deserialize(reader)?;
|
||||
let kind = match kind {
|
||||
1 => DictionaryKind::Fst,
|
||||
2 => DictionaryKind::SSTable,
|
||||
_ => {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::Other,
|
||||
format!("invalid dictionary kind: {kind}"),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
Ok(DictionaryFooter { kind, version })
|
||||
}
|
||||
}
|
||||
@@ -5,7 +5,7 @@ use std::{fmt, io};
|
||||
use async_trait::async_trait;
|
||||
use ownedbytes::{OwnedBytes, StableDeref};
|
||||
|
||||
use crate::{ByteCount, HasLen};
|
||||
use crate::HasLen;
|
||||
|
||||
/// Objects that represents files sections in tantivy.
|
||||
///
|
||||
@@ -216,11 +216,6 @@ impl FileSlice {
|
||||
pub fn slice_to(&self, to_offset: usize) -> FileSlice {
|
||||
self.slice(0..to_offset)
|
||||
}
|
||||
|
||||
/// Returns the byte count of the FileSlice.
|
||||
pub fn num_bytes(&self) -> ByteCount {
|
||||
self.range.len().into()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -5,18 +5,12 @@ use std::ops::Deref;
|
||||
pub use byteorder::LittleEndian as Endianness;
|
||||
|
||||
mod bitset;
|
||||
mod byte_count;
|
||||
mod datetime;
|
||||
mod dictionary_footer;
|
||||
pub mod file_slice;
|
||||
mod group_by;
|
||||
mod serialize;
|
||||
mod vint;
|
||||
mod writer;
|
||||
pub use bitset::*;
|
||||
pub use byte_count::ByteCount;
|
||||
pub use datetime::{DatePrecision, DateTime};
|
||||
pub use dictionary_footer::*;
|
||||
pub use group_by::GroupByIteratorExtended;
|
||||
pub use ownedbytes::{OwnedBytes, StableDeref};
|
||||
pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize};
|
||||
@@ -113,21 +107,6 @@ pub fn u64_to_f64(val: u64) -> f64 {
|
||||
})
|
||||
}
|
||||
|
||||
/// Replaces a given byte in the `bytes` slice of bytes.
|
||||
///
|
||||
/// This function assumes that the needle is rarely contained in the bytes string
|
||||
/// and offers a fast path if the needle is not present.
|
||||
pub fn replace_in_place(needle: u8, replacement: u8, bytes: &mut [u8]) {
|
||||
if !bytes.contains(&needle) {
|
||||
return;
|
||||
}
|
||||
for b in bytes {
|
||||
if *b == needle {
|
||||
*b = replacement;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test {
|
||||
|
||||
@@ -192,20 +171,4 @@ pub mod test {
|
||||
assert!(f64_to_u64(-2.0) < f64_to_u64(1.0));
|
||||
assert!(f64_to_u64(-2.0) < f64_to_u64(-1.5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_replace_in_place() {
|
||||
let test_aux = |before_replacement: &[u8], expected: &[u8]| {
|
||||
let mut bytes: Vec<u8> = before_replacement.to_vec();
|
||||
super::replace_in_place(b'b', b'c', &mut bytes);
|
||||
assert_eq!(&bytes[..], expected);
|
||||
};
|
||||
test_aux(b"", b"");
|
||||
test_aux(b"b", b"c");
|
||||
test_aux(b"baaa", b"caaa");
|
||||
test_aux(b"aaab", b"aaac");
|
||||
test_aux(b"aaabaa", b"aaacaa");
|
||||
test_aux(b"aaaaaa", b"aaaaaa");
|
||||
test_aux(b"bbbb", b"cccc");
|
||||
}
|
||||
}
|
||||
|
||||
130
examples-disabled/aggregation.rs
Normal file
130
examples-disabled/aggregation.rs
Normal file
@@ -0,0 +1,130 @@
|
||||
// # Aggregation example
|
||||
//
|
||||
// This example shows how you can use built-in aggregations.
|
||||
// We will use range buckets and compute the average in each bucket.
|
||||
//
|
||||
|
||||
use serde_json::Value;
|
||||
use tantivy::aggregation::agg_req::{
|
||||
Aggregation, Aggregations, BucketAggregation, BucketAggregationType, MetricAggregation,
|
||||
RangeAggregation,
|
||||
};
|
||||
use tantivy::aggregation::agg_result::AggregationResults;
|
||||
use tantivy::aggregation::metric::AverageAggregation;
|
||||
use tantivy::aggregation::AggregationCollector;
|
||||
use tantivy::query::TermQuery;
|
||||
use tantivy::schema::{self, IndexRecordOption, Schema, TextFieldIndexing};
|
||||
use tantivy::{doc, Index, Term};
|
||||
|
||||
fn main() -> tantivy::Result<()> {
|
||||
let mut schema_builder = Schema::builder();
|
||||
let text_fieldtype = schema::TextOptions::default()
|
||||
.set_indexing_options(
|
||||
TextFieldIndexing::default().set_index_option(IndexRecordOption::WithFreqs),
|
||||
)
|
||||
.set_stored();
|
||||
let text_field = schema_builder.add_text_field("text", text_fieldtype);
|
||||
let score_fieldtype =
|
||||
crate::schema::NumericOptions::default().set_fast();
|
||||
let highscore_field = schema_builder.add_f64_field("highscore", score_fieldtype.clone());
|
||||
let price_field = schema_builder.add_f64_field("price", score_fieldtype);
|
||||
|
||||
let schema = schema_builder.build();
|
||||
|
||||
// # Indexing documents
|
||||
//
|
||||
// Lets index a bunch of documents for this example.
|
||||
let index = Index::create_in_ram(schema);
|
||||
|
||||
let mut index_writer = index.writer(50_000_000)?;
|
||||
// writing the segment
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "cool",
|
||||
highscore_field => 1f64,
|
||||
price_field => 0f64,
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "cool",
|
||||
highscore_field => 3f64,
|
||||
price_field => 1f64,
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "cool",
|
||||
highscore_field => 5f64,
|
||||
price_field => 1f64,
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "nohit",
|
||||
highscore_field => 6f64,
|
||||
price_field => 2f64,
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "cool",
|
||||
highscore_field => 7f64,
|
||||
price_field => 2f64,
|
||||
))?;
|
||||
index_writer.commit()?;
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "cool",
|
||||
highscore_field => 11f64,
|
||||
price_field => 10f64,
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "cool",
|
||||
highscore_field => 14f64,
|
||||
price_field => 15f64,
|
||||
))?;
|
||||
|
||||
index_writer.add_document(doc!(
|
||||
text_field => "cool",
|
||||
highscore_field => 15f64,
|
||||
price_field => 20f64,
|
||||
))?;
|
||||
|
||||
index_writer.commit()?;
|
||||
|
||||
let reader = index.reader()?;
|
||||
let text_field = reader.searcher().schema().get_field("text").unwrap();
|
||||
|
||||
let term_query = TermQuery::new(
|
||||
Term::from_field_text(text_field, "cool"),
|
||||
IndexRecordOption::Basic,
|
||||
);
|
||||
|
||||
let sub_agg_req_1: Aggregations = vec![(
|
||||
"average_price".to_string(),
|
||||
Aggregation::Metric(MetricAggregation::Average(
|
||||
AverageAggregation::from_field_name("price".to_string()),
|
||||
)),
|
||||
)]
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
let agg_req_1: Aggregations = vec![(
|
||||
"score_ranges".to_string(),
|
||||
Aggregation::Bucket(BucketAggregation {
|
||||
bucket_agg: BucketAggregationType::Range(RangeAggregation {
|
||||
field: "highscore".to_string(),
|
||||
ranges: vec![
|
||||
(-1f64..9f64).into(),
|
||||
(9f64..14f64).into(),
|
||||
(14f64..20f64).into(),
|
||||
],
|
||||
..Default::default()
|
||||
}),
|
||||
sub_aggregation: sub_agg_req_1,
|
||||
}),
|
||||
)]
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap();
|
||||
|
||||
let res: Value = serde_json::to_value(agg_res)?;
|
||||
println!("{}", serde_json::to_string_pretty(&res)?);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -7,7 +7,9 @@
|
||||
// Of course, you can have a look at the tantivy's built-in collectors
|
||||
// such as the `CountCollector` for more examples.
|
||||
|
||||
use columnar::Column;
|
||||
use std::sync::Arc;
|
||||
|
||||
use fastfield_codecs::Column;
|
||||
// ---
|
||||
// Importing tantivy...
|
||||
use tantivy::collector::{Collector, SegmentCollector};
|
||||
@@ -95,7 +97,7 @@ impl Collector for StatsCollector {
|
||||
}
|
||||
|
||||
struct StatsSegmentCollector {
|
||||
fast_field_reader: Column,
|
||||
fast_field_reader: Arc<dyn Column<u64>>,
|
||||
stats: Stats,
|
||||
}
|
||||
|
||||
@@ -103,14 +105,10 @@ impl SegmentCollector for StatsSegmentCollector {
|
||||
type Fruit = Option<Stats>;
|
||||
|
||||
fn collect(&mut self, doc: u32, _score: Score) {
|
||||
// Since we know the values are single value, we could call `first_or_default_col` on the
|
||||
// column and fetch single values.
|
||||
for value in self.fast_field_reader.values_for_doc(doc) {
|
||||
let value = value as f64;
|
||||
self.stats.count += 1;
|
||||
self.stats.sum += value;
|
||||
self.stats.squared_sum += value * value;
|
||||
}
|
||||
let value = self.fast_field_reader.get_val(doc) as f64;
|
||||
self.stats.count += 1;
|
||||
self.stats.sum += value;
|
||||
self.stats.squared_sum += value * value;
|
||||
}
|
||||
|
||||
fn harvest(self) -> <Self as SegmentCollector>::Fruit {
|
||||
@@ -171,7 +169,7 @@ fn main() -> tantivy::Result<()> {
|
||||
let searcher = reader.searcher();
|
||||
let query_parser = QueryParser::for_index(&index, vec![product_name, product_description]);
|
||||
|
||||
// here we want to search for `broom` and use `StatsCollector` on the hits.
|
||||
// here we want to get a hit on the 'ken' in Frankenstein
|
||||
let query = query_parser.parse_query("broom")?;
|
||||
if let Some(stats) =
|
||||
searcher.search(&query, &StatsCollector::with_field("price".to_string()))?
|
||||
@@ -1,7 +1,7 @@
|
||||
// # Defining a tokenizer pipeline
|
||||
//
|
||||
// In this example, we'll see how to define a tokenizer
|
||||
// by creating a custom `NgramTokenizer`.
|
||||
// In this example, we'll see how to define a tokenizer pipeline
|
||||
// by aligning a bunch of `TokenFilter`.
|
||||
use tantivy::collector::TopDocs;
|
||||
use tantivy::query::QueryParser;
|
||||
use tantivy::schema::*;
|
||||
@@ -14,7 +14,6 @@ fn main() -> tantivy::Result<()> {
|
||||
.set_stored()
|
||||
.set_fast()
|
||||
.set_precision(tantivy::DatePrecision::Seconds);
|
||||
// Add `occurred_at` date field type
|
||||
let occurred_at = schema_builder.add_date_field("occurred_at", opts);
|
||||
let event_type = schema_builder.add_text_field("event", STRING | STORED);
|
||||
let schema = schema_builder.build();
|
||||
@@ -23,7 +22,6 @@ fn main() -> tantivy::Result<()> {
|
||||
let index = Index::create_in_ram(schema.clone());
|
||||
|
||||
let mut index_writer = index.writer(50_000_000)?;
|
||||
// The dates are passed as string in the RFC3339 format
|
||||
let doc = schema.parse_document(
|
||||
r#"{
|
||||
"occurred_at": "2022-06-22T12:53:50.53Z",
|
||||
@@ -43,16 +41,14 @@ fn main() -> tantivy::Result<()> {
|
||||
let reader = index.reader()?;
|
||||
let searcher = reader.searcher();
|
||||
|
||||
// # Search
|
||||
// # Default fields: event_type
|
||||
let query_parser = QueryParser::for_index(&index, vec![event_type]);
|
||||
{
|
||||
// Simple exact search on the date
|
||||
let query = query_parser.parse_query("occurred_at:\"2022-06-22T12:53:50.53Z\"")?;
|
||||
let query = query_parser.parse_query("event:comment")?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(5))?;
|
||||
assert_eq!(count_docs.len(), 1);
|
||||
}
|
||||
{
|
||||
// Range query on the date field
|
||||
let query = query_parser
|
||||
.parse_query(r#"occurred_at:[2022-06-22T12:58:00Z TO 2022-06-23T00:00:00Z}"#)?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(4))?;
|
||||
@@ -71,7 +71,7 @@ fn main() -> tantivy::Result<()> {
|
||||
let reader = index.reader()?;
|
||||
let searcher = reader.searcher();
|
||||
{
|
||||
let mut facet_collector = FacetCollector::for_field("classification");
|
||||
let mut facet_collector = FacetCollector::for_field(classification);
|
||||
facet_collector.add_facet("/Felidae");
|
||||
let facet_counts = searcher.search(&AllQuery, &facet_collector)?;
|
||||
// This lists all of the facet counts, right below "/Felidae".
|
||||
@@ -97,7 +97,7 @@ fn main() -> tantivy::Result<()> {
|
||||
let facet = Facet::from("/Felidae/Pantherinae");
|
||||
let facet_term = Term::from_facet(classification, &facet);
|
||||
let facet_term_query = TermQuery::new(facet_term, IndexRecordOption::Basic);
|
||||
let mut facet_collector = FacetCollector::for_field("classification");
|
||||
let mut facet_collector = FacetCollector::for_field(classification);
|
||||
facet_collector.add_facet("/Felidae/Pantherinae");
|
||||
let facet_counts = searcher.search(&facet_term_query, &facet_collector)?;
|
||||
let facets: Vec<(&Facet, u64)> = facet_counts.get("/Felidae/Pantherinae").collect();
|
||||
@@ -1,12 +1,3 @@
|
||||
// # Faceted Search With Tweak Score
|
||||
//
|
||||
// This example covers the faceted search functionalities of
|
||||
// tantivy.
|
||||
//
|
||||
// We will :
|
||||
// - define a text field "name" in our schema
|
||||
// - define a facet field "classification" in our schema
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use tantivy::collector::TopDocs;
|
||||
@@ -64,9 +55,8 @@ fn main() -> tantivy::Result<()> {
|
||||
.collect(),
|
||||
);
|
||||
let top_docs_by_custom_score =
|
||||
// Call TopDocs with a custom tweak score
|
||||
TopDocs::with_limit(2).tweak_score(move |segment_reader: &SegmentReader| {
|
||||
let ingredient_reader = segment_reader.facet_reader("ingredient").unwrap();
|
||||
let ingredient_reader = segment_reader.facet_reader(ingredient).unwrap();
|
||||
let facet_dict = ingredient_reader.facet_dict();
|
||||
|
||||
let query_ords: HashSet<u64> = facets
|
||||
@@ -74,10 +64,12 @@ fn main() -> tantivy::Result<()> {
|
||||
.filter_map(|key| facet_dict.term_ord(key.encoded_str()).unwrap())
|
||||
.collect();
|
||||
|
||||
let mut facet_ords_buffer: Vec<u64> = Vec::with_capacity(20);
|
||||
|
||||
move |doc: DocId, original_score: Score| {
|
||||
// Update the original score with a tweaked score
|
||||
let missing_ingredients = ingredient_reader
|
||||
.facet_ords(doc)
|
||||
ingredient_reader.facet_ords(doc, &mut facet_ords_buffer);
|
||||
let missing_ingredients = facet_ords_buffer
|
||||
.iter()
|
||||
.filter(|ord| !query_ords.contains(ord))
|
||||
.count();
|
||||
let tweak = 1.0 / 4_f32.powi(missing_ingredients as i32);
|
||||
73
examples-disabled/ip_field.rs
Normal file
73
examples-disabled/ip_field.rs
Normal file
@@ -0,0 +1,73 @@
|
||||
// # IP Address example
|
||||
//
|
||||
// This example shows how the ip field can be used
|
||||
// with IpV6 and IpV4.
|
||||
|
||||
use tantivy::collector::{Count, TopDocs};
|
||||
use tantivy::query::QueryParser;
|
||||
use tantivy::schema::{Schema, FAST, INDEXED, STORED, STRING};
|
||||
use tantivy::Index;
|
||||
|
||||
fn main() -> tantivy::Result<()> {
|
||||
// # Defining the schema
|
||||
let mut schema_builder = Schema::builder();
|
||||
let event_type = schema_builder.add_text_field("event_type", STRING | STORED);
|
||||
let ip = schema_builder.add_ip_addr_field("ip", STORED | INDEXED | FAST);
|
||||
let schema = schema_builder.build();
|
||||
|
||||
// # Indexing documents
|
||||
let index = Index::create_in_ram(schema.clone());
|
||||
|
||||
let mut index_writer = index.writer(50_000_000)?;
|
||||
let doc = schema.parse_document(
|
||||
r#"{
|
||||
"ip": "192.168.0.33",
|
||||
"event_type": "login"
|
||||
}"#,
|
||||
)?;
|
||||
index_writer.add_document(doc)?;
|
||||
let doc = schema.parse_document(
|
||||
r#"{
|
||||
"ip": "192.168.0.80",
|
||||
"event_type": "checkout"
|
||||
}"#,
|
||||
)?;
|
||||
index_writer.add_document(doc)?;
|
||||
let doc = schema.parse_document(
|
||||
r#"{
|
||||
"ip": "2001:0db8:85a3:0000:0000:8a2e:0370:7334",
|
||||
"event_type": "checkout"
|
||||
}"#,
|
||||
)?;
|
||||
|
||||
index_writer.add_document(doc)?;
|
||||
index_writer.commit()?;
|
||||
|
||||
let reader = index.reader()?;
|
||||
let searcher = reader.searcher();
|
||||
|
||||
let query_parser = QueryParser::for_index(&index, vec![event_type, ip]);
|
||||
{
|
||||
let query = query_parser.parse_query("ip:[192.168.0.0 TO 192.168.0.100]")?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(5))?;
|
||||
assert_eq!(count_docs.len(), 2);
|
||||
}
|
||||
{
|
||||
let query = query_parser.parse_query("ip:[192.168.1.0 TO 192.168.1.100]")?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(2))?;
|
||||
assert_eq!(count_docs.len(), 0);
|
||||
}
|
||||
{
|
||||
let query = query_parser.parse_query("ip:192.168.0.80")?;
|
||||
let count_docs = searcher.search(&*query, &Count)?;
|
||||
assert_eq!(count_docs, 1);
|
||||
}
|
||||
{
|
||||
// IpV6 needs to be escaped because it contains `:`
|
||||
let query = query_parser.parse_query("ip:\"2001:0db8:85a3:0000:0000:8a2e:0370:7334\"")?;
|
||||
let count_docs = searcher.search(&*query, &Count)?;
|
||||
assert_eq!(count_docs, 1);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -12,7 +12,7 @@
|
||||
use tantivy::collector::{Count, TopDocs};
|
||||
use tantivy::query::TermQuery;
|
||||
use tantivy::schema::*;
|
||||
use tantivy::tokenizer::{PreTokenizedString, SimpleTokenizer, Token, TokenStream, Tokenizer};
|
||||
use tantivy::tokenizer::{PreTokenizedString, SimpleTokenizer, Token, Tokenizer};
|
||||
use tantivy::{doc, Index, ReloadPolicy};
|
||||
use tempfile::TempDir;
|
||||
|
||||
@@ -50,13 +50,12 @@ fn main() -> tantivy::Result<()> {
|
||||
|
||||
// This tokenizer lowers all of the text (to help with stop word matching)
|
||||
// then removes all instances of `the` and `and` from the corpus
|
||||
let tokenizer = TextAnalyzer::builder(SimpleTokenizer)
|
||||
let tokenizer = TextAnalyzer::from(SimpleTokenizer)
|
||||
.filter(LowerCaser)
|
||||
.filter(StopWordFilter::remove(vec![
|
||||
"the".to_string(),
|
||||
"and".to_string(),
|
||||
]))
|
||||
.build();
|
||||
]));
|
||||
|
||||
index.tokenizers().register("stoppy", tokenizer);
|
||||
|
||||
@@ -17,6 +17,7 @@ use tantivy::{
|
||||
|
||||
type ProductId = u64;
|
||||
|
||||
/// Price
|
||||
type Price = u32;
|
||||
|
||||
pub trait PriceFetcher: Send + Sync + 'static {
|
||||
@@ -47,10 +48,7 @@ impl Warmer for DynamicPriceColumn {
|
||||
fn warm(&self, searcher: &Searcher) -> tantivy::Result<()> {
|
||||
for segment in searcher.segment_readers() {
|
||||
let key = (segment.segment_id(), segment.delete_opstamp());
|
||||
let product_id_reader = segment
|
||||
.fast_fields()
|
||||
.u64(&self.field)?
|
||||
.first_or_default_col(0);
|
||||
let product_id_reader = segment.fast_fields().u64(&self.field)?;
|
||||
let product_ids: Vec<ProductId> = segment
|
||||
.doc_ids_alive()
|
||||
.map(|doc| product_id_reader.get_val(doc))
|
||||
@@ -89,10 +87,10 @@ impl Warmer for DynamicPriceColumn {
|
||||
}
|
||||
}
|
||||
|
||||
// For the sake of this example, the table is just an editable HashMap behind a RwLock.
|
||||
// This map represents a map (ProductId -> Price)
|
||||
//
|
||||
// In practise, it could be fetching things from an external service, like a SQL table.
|
||||
/// For the sake of this example, the table is just an editable HashMap behind a RwLock.
|
||||
/// This map represents a map (ProductId -> Price)
|
||||
///
|
||||
/// In practise, it could be fetching things from an external service, like a SQL table.
|
||||
#[derive(Default, Clone)]
|
||||
pub struct ExternalPriceTable {
|
||||
prices: Arc<RwLock<HashMap<ProductId, Price>>>,
|
||||
@@ -1,319 +0,0 @@
|
||||
// # Aggregation example
|
||||
//
|
||||
// This example shows how you can use built-in aggregations.
|
||||
// We will use nested aggregations with buckets and metrics:
|
||||
// - Range buckets and compute the average in each bucket.
|
||||
// - Term aggregation and compute the min price in each bucket
|
||||
// ---
|
||||
|
||||
use serde_json::{Deserializer, Value};
|
||||
use tantivy::aggregation::agg_req::{
|
||||
Aggregation, Aggregations, BucketAggregation, BucketAggregationType, MetricAggregation,
|
||||
RangeAggregation,
|
||||
};
|
||||
use tantivy::aggregation::agg_result::AggregationResults;
|
||||
use tantivy::aggregation::bucket::RangeAggregationRange;
|
||||
use tantivy::aggregation::metric::AverageAggregation;
|
||||
use tantivy::aggregation::AggregationCollector;
|
||||
use tantivy::query::AllQuery;
|
||||
use tantivy::schema::{self, IndexRecordOption, Schema, TextFieldIndexing, FAST};
|
||||
use tantivy::Index;
|
||||
|
||||
fn main() -> tantivy::Result<()> {
|
||||
// # Create Schema
|
||||
//
|
||||
// Lets create a schema for a footwear shop, with 4 fields: name, category, stock and price.
|
||||
// category, stock and price will be fast fields as that's the requirement
|
||||
// for aggregation queries.
|
||||
//
|
||||
|
||||
let mut schema_builder = Schema::builder();
|
||||
// In preparation of the `TermsAggregation`, the category field is configured with:
|
||||
// - `set_fast`
|
||||
// - `raw` tokenizer
|
||||
//
|
||||
// The tokenizer is set to "raw", because the fast field uses the same dictionary as the
|
||||
// inverted index. (This behaviour will change in tantivy 0.20, where the fast field will
|
||||
// always be raw tokenized independent from the regular tokenizing)
|
||||
//
|
||||
let text_fieldtype = schema::TextOptions::default()
|
||||
.set_indexing_options(
|
||||
TextFieldIndexing::default()
|
||||
.set_index_option(IndexRecordOption::WithFreqs)
|
||||
.set_tokenizer("raw"),
|
||||
)
|
||||
.set_fast()
|
||||
.set_stored();
|
||||
schema_builder.add_text_field("category", text_fieldtype);
|
||||
schema_builder.add_f64_field("stock", FAST);
|
||||
schema_builder.add_f64_field("price", FAST);
|
||||
|
||||
let schema = schema_builder.build();
|
||||
|
||||
// # Indexing documents
|
||||
//
|
||||
// Lets index a bunch of documents for this example.
|
||||
let index = Index::create_in_ram(schema.clone());
|
||||
|
||||
let data = r#"{
|
||||
"name": "Almond Toe Court Shoes, Patent Black",
|
||||
"category": "Womens Footwear",
|
||||
"price": 99.00,
|
||||
"stock": 5
|
||||
}
|
||||
{
|
||||
"name": "Suede Shoes, Blue",
|
||||
"category": "Womens Footwear",
|
||||
"price": 42.00,
|
||||
"stock": 4
|
||||
}
|
||||
{
|
||||
"name": "Leather Driver Saddle Loafers, Tan",
|
||||
"category": "Mens Footwear",
|
||||
"price": 34.00,
|
||||
"stock": 12
|
||||
}
|
||||
{
|
||||
"name": "Flip Flops, Red",
|
||||
"category": "Mens Footwear",
|
||||
"price": 19.00,
|
||||
"stock": 6
|
||||
}
|
||||
{
|
||||
"name": "Flip Flops, Blue",
|
||||
"category": "Mens Footwear",
|
||||
"price": 19.00,
|
||||
"stock": 0
|
||||
}
|
||||
{
|
||||
"name": "Gold Button Cardigan, Black",
|
||||
"category": "Womens Casualwear",
|
||||
"price": 167.00,
|
||||
"stock": 6
|
||||
}
|
||||
{
|
||||
"name": "Cotton Shorts, Medium Red",
|
||||
"category": "Womens Casualwear",
|
||||
"price": 30.00,
|
||||
"stock": 5
|
||||
}
|
||||
{
|
||||
"name": "Fine Stripe Short SleeveShirt, Grey",
|
||||
"category": "Mens Casualwear",
|
||||
"price": 49.99,
|
||||
"stock": 9
|
||||
}
|
||||
{
|
||||
"name": "Fine Stripe Short SleeveShirt, Green",
|
||||
"category": "Mens Casualwear",
|
||||
"price": 49.99,
|
||||
"offer": 39.99,
|
||||
"stock": 9
|
||||
}
|
||||
{
|
||||
"name": "Sharkskin Waistcoat, Charcoal",
|
||||
"category": "Mens Formalwear",
|
||||
"price": 75.00,
|
||||
"stock": 2
|
||||
}
|
||||
{
|
||||
"name": "Lightweight Patch PocketBlazer, Deer",
|
||||
"category": "Mens Formalwear",
|
||||
"price": 175.50,
|
||||
"stock": 1
|
||||
}
|
||||
{
|
||||
"name": "Bird Print Dress, Black",
|
||||
"category": "Womens Formalwear",
|
||||
"price": 270.00,
|
||||
"stock": 10
|
||||
}
|
||||
{
|
||||
"name": "Mid Twist Cut-Out Dress, Pink",
|
||||
"category": "Womens Formalwear",
|
||||
"price": 540.00,
|
||||
"stock": 5
|
||||
}"#;
|
||||
|
||||
let stream = Deserializer::from_str(data).into_iter::<Value>();
|
||||
|
||||
let mut index_writer = index.writer(50_000_000)?;
|
||||
let mut num_indexed = 0;
|
||||
for value in stream {
|
||||
let doc = schema.parse_document(&serde_json::to_string(&value.unwrap())?)?;
|
||||
index_writer.add_document(doc)?;
|
||||
num_indexed += 1;
|
||||
if num_indexed > 4 {
|
||||
// Writing the first segment
|
||||
index_writer.commit()?;
|
||||
}
|
||||
}
|
||||
|
||||
// Writing the second segment
|
||||
index_writer.commit()?;
|
||||
|
||||
// We have two segments now. The `AggregationCollector` will run the aggregation on each
|
||||
// segment and then merge the results into an `IntermediateAggregationResult`.
|
||||
|
||||
let reader = index.reader()?;
|
||||
let searcher = reader.searcher();
|
||||
// ---
|
||||
// # Aggregation Query
|
||||
//
|
||||
//
|
||||
// We can construct the query by building the request structure or by deserializing from JSON.
|
||||
// The JSON API is more stable and therefore recommended.
|
||||
//
|
||||
// ## Request 1
|
||||
|
||||
let agg_req_str = r#"
|
||||
{
|
||||
"group_by_stock": {
|
||||
"aggs": {
|
||||
"average_price": { "avg": { "field": "price" } }
|
||||
},
|
||||
"range": {
|
||||
"field": "stock",
|
||||
"ranges": [
|
||||
{ "key": "few", "to": 1.0 },
|
||||
{ "key": "some", "from": 1.0, "to": 10.0 },
|
||||
{ "key": "many", "from": 10.0 }
|
||||
]
|
||||
}
|
||||
}
|
||||
} "#;
|
||||
|
||||
// In this Aggregation we want to get the average price for different groups, depending on how
|
||||
// many items are in stock. We define custom ranges `few`, `some`, `many` via the
|
||||
// range aggregation.
|
||||
// For every bucket we want the average price, so we create a nested metric aggregation on the
|
||||
// range bucket aggregation. Only buckets support nested aggregations.
|
||||
// ### Request JSON API
|
||||
//
|
||||
|
||||
let agg_req: Aggregations = serde_json::from_str(agg_req_str)?;
|
||||
let collector = AggregationCollector::from_aggs(agg_req, Default::default());
|
||||
|
||||
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap();
|
||||
let res2: Value = serde_json::to_value(agg_res)?;
|
||||
|
||||
// ### Request Rust API
|
||||
//
|
||||
// This is exactly the same request as above, but via the rust structures.
|
||||
//
|
||||
|
||||
let agg_req: Aggregations = vec![(
|
||||
"group_by_stock".to_string(),
|
||||
Aggregation::Bucket(Box::new(BucketAggregation {
|
||||
bucket_agg: BucketAggregationType::Range(RangeAggregation {
|
||||
field: "stock".to_string(),
|
||||
ranges: vec![
|
||||
RangeAggregationRange {
|
||||
key: Some("few".into()),
|
||||
from: None,
|
||||
to: Some(1f64),
|
||||
},
|
||||
RangeAggregationRange {
|
||||
key: Some("some".into()),
|
||||
from: Some(1f64),
|
||||
to: Some(10f64),
|
||||
},
|
||||
RangeAggregationRange {
|
||||
key: Some("many".into()),
|
||||
from: Some(10f64),
|
||||
to: None,
|
||||
},
|
||||
],
|
||||
..Default::default()
|
||||
}),
|
||||
sub_aggregation: vec![(
|
||||
"average_price".to_string(),
|
||||
Aggregation::Metric(MetricAggregation::Average(
|
||||
AverageAggregation::from_field_name("price".to_string()),
|
||||
)),
|
||||
)]
|
||||
.into_iter()
|
||||
.collect(),
|
||||
})),
|
||||
)]
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
let collector = AggregationCollector::from_aggs(agg_req, Default::default());
|
||||
// We use the `AllQuery` which will pass all documents to the AggregationCollector.
|
||||
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap();
|
||||
|
||||
let res1: Value = serde_json::to_value(agg_res)?;
|
||||
|
||||
// ### Aggregation Result
|
||||
//
|
||||
// The resulting structure deserializes in the same JSON format as elastic search.
|
||||
//
|
||||
let expected_res = r#"
|
||||
{
|
||||
"group_by_stock":{
|
||||
"buckets":[
|
||||
{"average_price":{"value":19.0},"doc_count":1,"key":"few","to":1.0},
|
||||
{"average_price":{"value":124.748},"doc_count":10,"from":1.0,"key":"some","to":10.0},
|
||||
{"average_price":{"value":152.0},"doc_count":2,"from":10.0,"key":"many"}
|
||||
]
|
||||
}
|
||||
}
|
||||
"#;
|
||||
let expected_json: Value = serde_json::from_str(expected_res)?;
|
||||
assert_eq!(expected_json, res1);
|
||||
assert_eq!(expected_json, res2);
|
||||
|
||||
// ### Request 2
|
||||
//
|
||||
// Now we are interested in the minimum price per category, so we create a bucket per
|
||||
// category via `TermsAggregation`. We are interested in the highest minimum prices, and set the
|
||||
// order of the buckets `"order": { "min_price": "desc" }` to be sorted by the the metric of
|
||||
// the sub aggregation. (awesome)
|
||||
//
|
||||
let agg_req_str = r#"
|
||||
{
|
||||
"min_price_per_category": {
|
||||
"aggs": {
|
||||
"min_price": { "min": { "field": "price" } }
|
||||
},
|
||||
"terms": {
|
||||
"field": "category",
|
||||
"min_doc_count": 1,
|
||||
"order": { "min_price": "desc" }
|
||||
}
|
||||
}
|
||||
} "#;
|
||||
|
||||
let agg_req: Aggregations = serde_json::from_str(agg_req_str)?;
|
||||
|
||||
let collector = AggregationCollector::from_aggs(agg_req, Default::default());
|
||||
|
||||
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap();
|
||||
let res: Value = serde_json::to_value(agg_res)?;
|
||||
|
||||
// Minimum price per category, sorted by minimum price descending
|
||||
//
|
||||
// As you can see, the starting prices for `Formalwear` are higher than `Casualwear`.
|
||||
//
|
||||
let expected_res = r#"
|
||||
{
|
||||
"min_price_per_category": {
|
||||
"buckets": [
|
||||
{ "doc_count": 2, "key": "Womens Formalwear", "min_price": { "value": 270.0 } },
|
||||
{ "doc_count": 2, "key": "Mens Formalwear", "min_price": { "value": 75.0 } },
|
||||
{ "doc_count": 2, "key": "Mens Casualwear", "min_price": { "value": 49.99 } },
|
||||
{ "doc_count": 2, "key": "Womens Footwear", "min_price": { "value": 42.0 } },
|
||||
{ "doc_count": 2, "key": "Womens Casualwear", "min_price": { "value": 30.0 } },
|
||||
{ "doc_count": 3, "key": "Mens Footwear", "min_price": { "value": 19.0 } }
|
||||
],
|
||||
"sum_other_doc_count": 0
|
||||
}
|
||||
}
|
||||
"#;
|
||||
let expected_json: Value = serde_json::from_str(expected_res)?;
|
||||
|
||||
assert_eq!(expected_json, res);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,167 +0,0 @@
|
||||
// # Basic Example
|
||||
//
|
||||
// This example covers the basic functionalities of
|
||||
// tantivy.
|
||||
//
|
||||
// We will :
|
||||
// - define our schema
|
||||
// - create an index in a directory
|
||||
// - index a few documents into our index
|
||||
// - search for the best document matching a basic query
|
||||
// - retrieve the best document's original content.
|
||||
// ---
|
||||
// Importing tantivy...
|
||||
use tantivy::collector::{Count, TopDocs};
|
||||
use tantivy::query::FuzzyTermQuery;
|
||||
use tantivy::schema::*;
|
||||
use tantivy::{doc, Index, ReloadPolicy};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn main() -> tantivy::Result<()> {
|
||||
// Let's create a temporary directory for the
|
||||
// sake of this example
|
||||
let index_path = TempDir::new()?;
|
||||
|
||||
// # Defining the schema
|
||||
//
|
||||
// The Tantivy index requires a very strict schema.
|
||||
// The schema declares which fields are in the index,
|
||||
// and for each field, its type and "the way it should
|
||||
// be indexed".
|
||||
|
||||
// First we need to define a schema ...
|
||||
let mut schema_builder = Schema::builder();
|
||||
|
||||
// Our first field is title.
|
||||
// We want full-text search for it, and we also want
|
||||
// to be able to retrieve the document after the search.
|
||||
//
|
||||
// `TEXT | STORED` is some syntactic sugar to describe
|
||||
// that.
|
||||
//
|
||||
// `TEXT` means the field should be tokenized and indexed,
|
||||
// along with its term frequency and term positions.
|
||||
//
|
||||
// `STORED` means that the field will also be saved
|
||||
// in a compressed, row-oriented key-value store.
|
||||
// This store is useful for reconstructing the
|
||||
// documents that were selected during the search phase.
|
||||
let title = schema_builder.add_text_field("title", TEXT | STORED);
|
||||
|
||||
let schema = schema_builder.build();
|
||||
|
||||
// # Indexing documents
|
||||
//
|
||||
// Let's create a brand new index.
|
||||
//
|
||||
// This will actually just save a meta.json
|
||||
// with our schema in the directory.
|
||||
let index = Index::create_in_dir(&index_path, schema.clone())?;
|
||||
|
||||
// To insert a document we will need an index writer.
|
||||
// There must be only one writer at a time.
|
||||
// This single `IndexWriter` is already
|
||||
// multithreaded.
|
||||
//
|
||||
// Here we give tantivy a budget of `50MB`.
|
||||
// Using a bigger memory_arena for the indexer may increase
|
||||
// throughput, but 50 MB is already plenty.
|
||||
let mut index_writer = index.writer(50_000_000)?;
|
||||
|
||||
// Let's index our documents!
|
||||
// We first need a handle on the title and the body field.
|
||||
|
||||
// ### Adding documents
|
||||
//
|
||||
index_writer.add_document(doc!(
|
||||
title => "The Name of the Wind",
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
title => "The Diary of Muadib",
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
title => "A Dairy Cow",
|
||||
))?;
|
||||
index_writer.add_document(doc!(
|
||||
title => "The Diary of a Young Girl",
|
||||
))?;
|
||||
index_writer.commit()?;
|
||||
|
||||
// ### Committing
|
||||
//
|
||||
// At this point our documents are not searchable.
|
||||
//
|
||||
//
|
||||
// We need to call `.commit()` explicitly to force the
|
||||
// `index_writer` to finish processing the documents in the queue,
|
||||
// flush the current index to the disk, and advertise
|
||||
// the existence of new documents.
|
||||
//
|
||||
// This call is blocking.
|
||||
index_writer.commit()?;
|
||||
|
||||
// If `.commit()` returns correctly, then all of the
|
||||
// documents that have been added are guaranteed to be
|
||||
// persistently indexed.
|
||||
//
|
||||
// In the scenario of a crash or a power failure,
|
||||
// tantivy behaves as if it has rolled back to its last
|
||||
// commit.
|
||||
|
||||
// # Searching
|
||||
//
|
||||
// ### Searcher
|
||||
//
|
||||
// A reader is required first in order to search an index.
|
||||
// It acts as a `Searcher` pool that reloads itself,
|
||||
// depending on a `ReloadPolicy`.
|
||||
//
|
||||
// For a search server you will typically create one reader for the entire lifetime of your
|
||||
// program, and acquire a new searcher for every single request.
|
||||
//
|
||||
// In the code below, we rely on the 'ON_COMMIT' policy: the reader
|
||||
// will reload the index automatically after each commit.
|
||||
let reader = index
|
||||
.reader_builder()
|
||||
.reload_policy(ReloadPolicy::OnCommit)
|
||||
.try_into()?;
|
||||
|
||||
// We now need to acquire a searcher.
|
||||
//
|
||||
// A searcher points to a snapshotted, immutable version of the index.
|
||||
//
|
||||
// Some search experience might require more than
|
||||
// one query. Using the same searcher ensures that all of these queries will run on the
|
||||
// same version of the index.
|
||||
//
|
||||
// Acquiring a `searcher` is very cheap.
|
||||
//
|
||||
// You should acquire a searcher every time you start processing a request and
|
||||
// and release it right after your query is finished.
|
||||
let searcher = reader.searcher();
|
||||
|
||||
// ### FuzzyTermQuery
|
||||
{
|
||||
let term = Term::from_field_text(title, "Diary");
|
||||
let query = FuzzyTermQuery::new(term, 2, true);
|
||||
|
||||
let (top_docs, count) = searcher
|
||||
.search(&query, &(TopDocs::with_limit(5), Count))
|
||||
.unwrap();
|
||||
assert_eq!(count, 3);
|
||||
assert_eq!(top_docs.len(), 3);
|
||||
for (score, doc_address) in top_docs {
|
||||
let retrieved_doc = searcher.doc(doc_address)?;
|
||||
// Note that the score is not lower for the fuzzy hit.
|
||||
// There's an issue open for that: https://github.com/quickwit-oss/tantivy/issues/563
|
||||
println!("score {score:?} doc {}", schema.to_json(&retrieved_doc));
|
||||
// score 1.0 doc {"title":["The Diary of Muadib"]}
|
||||
//
|
||||
// score 1.0 doc {"title":["The Diary of a Young Girl"]}
|
||||
//
|
||||
// score 1.0 doc {"title":["A Dairy Cow"]}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,107 +0,0 @@
|
||||
// # IP Address example
|
||||
//
|
||||
// This example shows how the ip field can be used
|
||||
// with IpV6 and IpV4.
|
||||
|
||||
use tantivy::collector::{Count, TopDocs};
|
||||
use tantivy::query::QueryParser;
|
||||
use tantivy::schema::{Schema, FAST, INDEXED, STORED, STRING};
|
||||
use tantivy::Index;
|
||||
|
||||
fn main() -> tantivy::Result<()> {
|
||||
// # Defining the schema
|
||||
// We set the IP field as `INDEXED`, so it can be searched
|
||||
// `FAST` will create a fast field. The fast field will be used to execute search queries.
|
||||
// `FAST` is not a requirement for range queries, it can also be executed on the inverted index
|
||||
// which is created by `INDEXED`.
|
||||
let mut schema_builder = Schema::builder();
|
||||
let event_type = schema_builder.add_text_field("event_type", STRING | STORED);
|
||||
let ip = schema_builder.add_ip_addr_field("ip", STORED | INDEXED | FAST);
|
||||
let schema = schema_builder.build();
|
||||
|
||||
// # Indexing documents
|
||||
let index = Index::create_in_ram(schema.clone());
|
||||
|
||||
let mut index_writer = index.writer(50_000_000)?;
|
||||
|
||||
// ### IPv4
|
||||
// Adding documents that contain an IPv4 address. Notice that the IP addresses are passed as
|
||||
// `String`. Since the field is of type ip, we parse the IP address from the string and store it
|
||||
// internally as IPv6.
|
||||
let doc = schema.parse_document(
|
||||
r#"{
|
||||
"ip": "192.168.0.33",
|
||||
"event_type": "login"
|
||||
}"#,
|
||||
)?;
|
||||
index_writer.add_document(doc)?;
|
||||
let doc = schema.parse_document(
|
||||
r#"{
|
||||
"ip": "192.168.0.80",
|
||||
"event_type": "checkout"
|
||||
}"#,
|
||||
)?;
|
||||
index_writer.add_document(doc)?;
|
||||
// ### IPv6
|
||||
// Adding a document that contains an IPv6 address.
|
||||
let doc = schema.parse_document(
|
||||
r#"{
|
||||
"ip": "2001:0db8:85a3:0000:0000:8a2e:0370:7334",
|
||||
"event_type": "checkout"
|
||||
}"#,
|
||||
)?;
|
||||
|
||||
index_writer.add_document(doc)?;
|
||||
// Commit will create a segment containing our documents.
|
||||
index_writer.commit()?;
|
||||
|
||||
let reader = index.reader()?;
|
||||
let searcher = reader.searcher();
|
||||
|
||||
// # Search
|
||||
// Range queries on IPv4. Since we created a fast field, the fast field will be used to execute
|
||||
// the search.
|
||||
// ### Range Queries
|
||||
let query_parser = QueryParser::for_index(&index, vec![event_type, ip]);
|
||||
{
|
||||
// Inclusive range queries
|
||||
let query = query_parser.parse_query("ip:[192.168.0.80 TO 192.168.0.100]")?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(5))?;
|
||||
assert_eq!(count_docs.len(), 1);
|
||||
}
|
||||
{
|
||||
// Exclusive range queries
|
||||
let query = query_parser.parse_query("ip:{192.168.0.80 TO 192.168.1.100]")?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(2))?;
|
||||
assert_eq!(count_docs.len(), 0);
|
||||
}
|
||||
{
|
||||
// Find docs with IP addresses smaller equal 192.168.1.100
|
||||
let query = query_parser.parse_query("ip:[* TO 192.168.1.100]")?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(2))?;
|
||||
assert_eq!(count_docs.len(), 2);
|
||||
}
|
||||
{
|
||||
// Find docs with IP addresses smaller than 192.168.1.100
|
||||
let query = query_parser.parse_query("ip:[* TO 192.168.1.100}")?;
|
||||
let count_docs = searcher.search(&*query, &TopDocs::with_limit(2))?;
|
||||
assert_eq!(count_docs.len(), 2);
|
||||
}
|
||||
|
||||
// ### Exact Queries
|
||||
// Exact search on IPv4.
|
||||
{
|
||||
let query = query_parser.parse_query("ip:192.168.0.80")?;
|
||||
let count_docs = searcher.search(&*query, &Count)?;
|
||||
assert_eq!(count_docs, 1);
|
||||
}
|
||||
// Exact search on IPv6.
|
||||
// IpV6 addresses need to be quoted because they contain `:`
|
||||
{
|
||||
let query = query_parser.parse_query("ip:\"2001:0db8:85a3:0000:0000:8a2e:0370:7334\"")?;
|
||||
let count_docs = searcher.search(&*query, &Count)?;
|
||||
assert_eq!(count_docs, 1);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,95 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::ByteCount;
|
||||
|
||||
use super::collector::DEFAULT_MEMORY_LIMIT;
|
||||
use super::{AggregationError, DEFAULT_BUCKET_LIMIT};
|
||||
use crate::TantivyError;
|
||||
|
||||
/// An estimate for memory consumption. Non recursive
|
||||
pub trait MemoryConsumption {
|
||||
fn memory_consumption(&self) -> usize;
|
||||
}
|
||||
|
||||
impl<K, V, S> MemoryConsumption for HashMap<K, V, S> {
|
||||
fn memory_consumption(&self) -> usize {
|
||||
let num_items = self.capacity();
|
||||
(std::mem::size_of::<K>() + std::mem::size_of::<V>()) * num_items
|
||||
}
|
||||
}
|
||||
|
||||
/// Aggregation memory limit after which the request fails. Defaults to DEFAULT_MEMORY_LIMIT
|
||||
/// (500MB). The limit is shared by all SegmentCollectors
|
||||
pub struct AggregationLimits {
|
||||
/// The counter which is shared between the aggregations for one request.
|
||||
memory_consumption: Arc<AtomicU64>,
|
||||
/// The memory_limit in bytes
|
||||
memory_limit: ByteCount,
|
||||
/// The maximum number of buckets _returned_
|
||||
/// This is not counting intermediate buckets.
|
||||
bucket_limit: u32,
|
||||
}
|
||||
impl Clone for AggregationLimits {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
memory_consumption: Arc::clone(&self.memory_consumption),
|
||||
memory_limit: self.memory_limit,
|
||||
bucket_limit: self.bucket_limit,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AggregationLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
memory_consumption: Default::default(),
|
||||
memory_limit: DEFAULT_MEMORY_LIMIT.into(),
|
||||
bucket_limit: DEFAULT_BUCKET_LIMIT,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AggregationLimits {
|
||||
/// *memory_limit*
|
||||
/// memory_limit is defined in bytes.
|
||||
/// Aggregation fails when the estimated memory consumption of the aggregation is higher than
|
||||
/// memory_limit.
|
||||
/// memory_limit will default to `DEFAULT_MEMORY_LIMIT` (500MB)
|
||||
///
|
||||
/// *bucket_limit*
|
||||
/// Limits the maximum number of buckets returned from an aggregation request.
|
||||
/// bucket_limit will default to `DEFAULT_BUCKET_LIMIT` (65000)
|
||||
pub fn new(memory_limit: Option<u64>, bucket_limit: Option<u32>) -> Self {
|
||||
Self {
|
||||
memory_consumption: Default::default(),
|
||||
memory_limit: memory_limit.unwrap_or(DEFAULT_MEMORY_LIMIT).into(),
|
||||
bucket_limit: bucket_limit.unwrap_or(DEFAULT_BUCKET_LIMIT),
|
||||
}
|
||||
}
|
||||
pub(crate) fn validate_memory_consumption(&self) -> crate::Result<()> {
|
||||
if self.get_memory_consumed() > self.memory_limit {
|
||||
return Err(TantivyError::AggregationError(
|
||||
AggregationError::MemoryExceeded {
|
||||
limit: self.memory_limit,
|
||||
current: self.get_memory_consumed(),
|
||||
},
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
pub(crate) fn add_memory_consumed(&self, num_bytes: u64) {
|
||||
self.memory_consumption
|
||||
.fetch_add(num_bytes, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
/// Returns the estimated memory consumed by the aggregations
|
||||
pub fn get_memory_consumed(&self) -> ByteCount {
|
||||
self.memory_consumption
|
||||
.load(std::sync::atomic::Ordering::Relaxed)
|
||||
.into()
|
||||
}
|
||||
pub(crate) fn get_bucket_limit(&self) -> u32 {
|
||||
self.bucket_limit
|
||||
}
|
||||
}
|
||||
@@ -16,14 +16,14 @@
|
||||
//! let agg_req1: Aggregations = vec![
|
||||
//! (
|
||||
//! "range".to_string(),
|
||||
//! Aggregation::Bucket(Box::new(BucketAggregation {
|
||||
//! Aggregation::Bucket(BucketAggregation {
|
||||
//! bucket_agg: BucketAggregationType::Range(RangeAggregation{
|
||||
//! field: "score".to_string(),
|
||||
//! ranges: vec![(3f64..7f64).into(), (7f64..20f64).into()],
|
||||
//! keyed: false,
|
||||
//! }),
|
||||
//! sub_aggregation: Default::default(),
|
||||
//! })),
|
||||
//! }),
|
||||
//! ),
|
||||
//! ]
|
||||
//! .into_iter()
|
||||
@@ -50,7 +50,7 @@ use std::collections::{HashMap, HashSet};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
pub use super::bucket::RangeAggregation;
|
||||
use super::bucket::{DateHistogramAggregationReq, HistogramAggregation, TermsAggregation};
|
||||
use super::bucket::{HistogramAggregation, TermsAggregation};
|
||||
use super::metric::{
|
||||
AverageAggregation, CountAggregation, MaxAggregation, MinAggregation, StatsAggregation,
|
||||
SumAggregation,
|
||||
@@ -110,13 +110,10 @@ impl BucketAggregationInternal {
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
pub(crate) fn as_histogram(&self) -> crate::Result<Option<HistogramAggregation>> {
|
||||
pub(crate) fn as_histogram(&self) -> Option<&HistogramAggregation> {
|
||||
match &self.bucket_agg {
|
||||
BucketAggregationType::Histogram(histogram) => Ok(Some(histogram.clone())),
|
||||
BucketAggregationType::DateHistogram(histogram) => {
|
||||
Ok(Some(histogram.to_histogram_req()?))
|
||||
}
|
||||
_ => Ok(None),
|
||||
BucketAggregationType::Histogram(histogram) => Some(histogram),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
pub(crate) fn as_term(&self) -> Option<&TermsAggregation> {
|
||||
@@ -127,6 +124,15 @@ impl BucketAggregationInternal {
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract all fields, where the term directory is used in the tree.
|
||||
pub fn get_term_dict_field_names(aggs: &Aggregations) -> HashSet<String> {
|
||||
let mut term_dict_field_names = Default::default();
|
||||
for el in aggs.values() {
|
||||
el.get_term_dict_field_names(&mut term_dict_field_names)
|
||||
}
|
||||
term_dict_field_names
|
||||
}
|
||||
|
||||
/// Extract all fast field names used in the tree.
|
||||
pub fn get_fast_field_names(aggs: &Aggregations) -> HashSet<String> {
|
||||
let mut fast_field_names = Default::default();
|
||||
@@ -143,18 +149,22 @@ pub fn get_fast_field_names(aggs: &Aggregations) -> HashSet<String> {
|
||||
#[serde(untagged)]
|
||||
pub enum Aggregation {
|
||||
/// Bucket aggregation, see [`BucketAggregation`] for details.
|
||||
Bucket(Box<BucketAggregation>),
|
||||
Bucket(BucketAggregation),
|
||||
/// Metric aggregation, see [`MetricAggregation`] for details.
|
||||
Metric(MetricAggregation),
|
||||
}
|
||||
|
||||
impl Aggregation {
|
||||
fn get_term_dict_field_names(&self, term_field_names: &mut HashSet<String>) {
|
||||
if let Aggregation::Bucket(bucket) = self {
|
||||
bucket.get_term_dict_field_names(term_field_names)
|
||||
}
|
||||
}
|
||||
|
||||
fn get_fast_field_names(&self, fast_field_names: &mut HashSet<String>) {
|
||||
match self {
|
||||
Aggregation::Bucket(bucket) => bucket.get_fast_field_names(fast_field_names),
|
||||
Aggregation::Metric(metric) => {
|
||||
fast_field_names.insert(metric.get_fast_field_name().to_string());
|
||||
}
|
||||
Aggregation::Metric(metric) => metric.get_fast_field_names(fast_field_names),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -183,9 +193,14 @@ pub struct BucketAggregation {
|
||||
}
|
||||
|
||||
impl BucketAggregation {
|
||||
fn get_term_dict_field_names(&self, term_dict_field_names: &mut HashSet<String>) {
|
||||
if let BucketAggregationType::Terms(terms) = &self.bucket_agg {
|
||||
term_dict_field_names.insert(terms.field.to_string());
|
||||
}
|
||||
term_dict_field_names.extend(get_term_dict_field_names(&self.sub_aggregation));
|
||||
}
|
||||
fn get_fast_field_names(&self, fast_field_names: &mut HashSet<String>) {
|
||||
let fast_field_name = self.bucket_agg.get_fast_field_name();
|
||||
fast_field_names.insert(fast_field_name.to_string());
|
||||
self.bucket_agg.get_fast_field_names(fast_field_names);
|
||||
fast_field_names.extend(get_fast_field_names(&self.sub_aggregation));
|
||||
}
|
||||
}
|
||||
@@ -199,22 +214,20 @@ pub enum BucketAggregationType {
|
||||
/// Put data into buckets of user-defined ranges.
|
||||
#[serde(rename = "histogram")]
|
||||
Histogram(HistogramAggregation),
|
||||
/// Put data into buckets of user-defined ranges.
|
||||
#[serde(rename = "date_histogram")]
|
||||
DateHistogram(DateHistogramAggregationReq),
|
||||
/// Put data into buckets of terms.
|
||||
#[serde(rename = "terms")]
|
||||
Terms(TermsAggregation),
|
||||
}
|
||||
|
||||
impl BucketAggregationType {
|
||||
fn get_fast_field_name(&self) -> &str {
|
||||
fn get_fast_field_names(&self, fast_field_names: &mut HashSet<String>) {
|
||||
match self {
|
||||
BucketAggregationType::Terms(terms) => terms.field.as_str(),
|
||||
BucketAggregationType::Range(range) => range.field.as_str(),
|
||||
BucketAggregationType::Histogram(histogram) => histogram.field.as_str(),
|
||||
BucketAggregationType::DateHistogram(histogram) => histogram.field.as_str(),
|
||||
}
|
||||
BucketAggregationType::Terms(terms) => fast_field_names.insert(terms.field.to_string()),
|
||||
BucketAggregationType::Range(range) => fast_field_names.insert(range.field.to_string()),
|
||||
BucketAggregationType::Histogram(histogram) => {
|
||||
fast_field_names.insert(histogram.field.to_string())
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
@@ -249,15 +262,16 @@ pub enum MetricAggregation {
|
||||
}
|
||||
|
||||
impl MetricAggregation {
|
||||
fn get_fast_field_name(&self) -> &str {
|
||||
match self {
|
||||
fn get_fast_field_names(&self, fast_field_names: &mut HashSet<String>) {
|
||||
let fast_field_name = match self {
|
||||
MetricAggregation::Average(avg) => avg.field_name(),
|
||||
MetricAggregation::Count(count) => count.field_name(),
|
||||
MetricAggregation::Max(max) => max.field_name(),
|
||||
MetricAggregation::Min(min) => min.field_name(),
|
||||
MetricAggregation::Stats(stats) => stats.field_name(),
|
||||
MetricAggregation::Sum(sum) => sum.field_name(),
|
||||
}
|
||||
};
|
||||
fast_field_names.insert(fast_field_name.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -301,7 +315,7 @@ mod tests {
|
||||
fn serialize_to_json_test() {
|
||||
let agg_req1: Aggregations = vec![(
|
||||
"range".to_string(),
|
||||
Aggregation::Bucket(Box::new(BucketAggregation {
|
||||
Aggregation::Bucket(BucketAggregation {
|
||||
bucket_agg: BucketAggregationType::Range(RangeAggregation {
|
||||
field: "score".to_string(),
|
||||
ranges: vec![
|
||||
@@ -313,7 +327,7 @@ mod tests {
|
||||
keyed: true,
|
||||
}),
|
||||
sub_aggregation: Default::default(),
|
||||
})),
|
||||
}),
|
||||
)]
|
||||
.into_iter()
|
||||
.collect();
|
||||
@@ -351,7 +365,7 @@ mod tests {
|
||||
let agg_req2: Aggregations = vec![
|
||||
(
|
||||
"range".to_string(),
|
||||
Aggregation::Bucket(Box::new(BucketAggregation {
|
||||
Aggregation::Bucket(BucketAggregation {
|
||||
bucket_agg: BucketAggregationType::Range(RangeAggregation {
|
||||
field: "score2".to_string(),
|
||||
ranges: vec![
|
||||
@@ -363,7 +377,7 @@ mod tests {
|
||||
..Default::default()
|
||||
}),
|
||||
sub_aggregation: Default::default(),
|
||||
})),
|
||||
}),
|
||||
),
|
||||
(
|
||||
"metric".to_string(),
|
||||
@@ -377,7 +391,7 @@ mod tests {
|
||||
|
||||
let agg_req1: Aggregations = vec![(
|
||||
"range".to_string(),
|
||||
Aggregation::Bucket(Box::new(BucketAggregation {
|
||||
Aggregation::Bucket(BucketAggregation {
|
||||
bucket_agg: BucketAggregationType::Range(RangeAggregation {
|
||||
field: "score".to_string(),
|
||||
ranges: vec![
|
||||
@@ -389,7 +403,7 @@ mod tests {
|
||||
..Default::default()
|
||||
}),
|
||||
sub_aggregation: agg_req2,
|
||||
})),
|
||||
}),
|
||||
)]
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
@@ -1,18 +1,22 @@
|
||||
//! This will enhance the request tree with access to the fastfield and metadata.
|
||||
|
||||
use columnar::{Column, ColumnBlockAccessor, ColumnType, StrColumn};
|
||||
use std::rc::Rc;
|
||||
use std::sync::atomic::AtomicU32;
|
||||
use std::sync::Arc;
|
||||
|
||||
use fastfield_codecs::Column;
|
||||
|
||||
use super::agg_req::{Aggregation, Aggregations, BucketAggregationType, MetricAggregation};
|
||||
use super::bucket::{
|
||||
DateHistogramAggregationReq, HistogramAggregation, RangeAggregation, TermsAggregation,
|
||||
};
|
||||
use super::bucket::{HistogramAggregation, RangeAggregation, TermsAggregation};
|
||||
use super::metric::{
|
||||
AverageAggregation, CountAggregation, MaxAggregation, MinAggregation, StatsAggregation,
|
||||
SumAggregation,
|
||||
};
|
||||
use super::segment_agg_result::AggregationLimits;
|
||||
use super::segment_agg_result::BucketCount;
|
||||
use super::VecWithNames;
|
||||
use crate::SegmentReader;
|
||||
use crate::fastfield::{type_and_cardinality, MultiValuedFastFieldReader};
|
||||
use crate::schema::Type;
|
||||
use crate::{InvertedIndexReader, SegmentReader, TantivyError};
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct AggregationsWithAccessor {
|
||||
@@ -33,17 +37,42 @@ impl AggregationsWithAccessor {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum FastFieldAccessor {
|
||||
Multi(MultiValuedFastFieldReader<u64>),
|
||||
Single(Arc<dyn Column<u64>>),
|
||||
}
|
||||
impl FastFieldAccessor {
|
||||
pub fn as_single(&self) -> Option<&dyn Column<u64>> {
|
||||
match self {
|
||||
FastFieldAccessor::Multi(_) => None,
|
||||
FastFieldAccessor::Single(reader) => Some(&**reader),
|
||||
}
|
||||
}
|
||||
pub fn into_single(self) -> Option<Arc<dyn Column<u64>>> {
|
||||
match self {
|
||||
FastFieldAccessor::Multi(_) => None,
|
||||
FastFieldAccessor::Single(reader) => Some(reader),
|
||||
}
|
||||
}
|
||||
pub fn as_multi(&self) -> Option<&MultiValuedFastFieldReader<u64>> {
|
||||
match self {
|
||||
FastFieldAccessor::Multi(reader) => Some(reader),
|
||||
FastFieldAccessor::Single(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct BucketAggregationWithAccessor {
|
||||
/// In general there can be buckets without fast field access, e.g. buckets that are created
|
||||
/// based on search terms. So eventually this needs to be Option or moved.
|
||||
pub(crate) accessor: Column<u64>,
|
||||
pub(crate) str_dict_column: Option<StrColumn>,
|
||||
pub(crate) field_type: ColumnType,
|
||||
pub(crate) accessor: FastFieldAccessor,
|
||||
pub(crate) inverted_index: Option<Arc<InvertedIndexReader>>,
|
||||
pub(crate) field_type: Type,
|
||||
pub(crate) bucket_agg: BucketAggregationType,
|
||||
pub(crate) sub_aggregation: AggregationsWithAccessor,
|
||||
pub(crate) limits: AggregationLimits,
|
||||
pub(crate) column_block_accessor: ColumnBlockAccessor<u64>,
|
||||
pub(crate) bucket_count: BucketCount,
|
||||
}
|
||||
|
||||
impl BucketAggregationWithAccessor {
|
||||
@@ -51,25 +80,23 @@ impl BucketAggregationWithAccessor {
|
||||
bucket: &BucketAggregationType,
|
||||
sub_aggregation: &Aggregations,
|
||||
reader: &SegmentReader,
|
||||
limits: AggregationLimits,
|
||||
bucket_count: Rc<AtomicU32>,
|
||||
max_bucket_count: u32,
|
||||
) -> crate::Result<BucketAggregationWithAccessor> {
|
||||
let mut str_dict_column = None;
|
||||
let mut inverted_index = None;
|
||||
let (accessor, field_type) = match &bucket {
|
||||
BucketAggregationType::Range(RangeAggregation {
|
||||
field: field_name, ..
|
||||
}) => get_ff_reader_and_validate(reader, field_name)?,
|
||||
}) => get_ff_reader_and_validate(reader, field_name, Cardinality::SingleValue)?,
|
||||
BucketAggregationType::Histogram(HistogramAggregation {
|
||||
field: field_name, ..
|
||||
}) => get_ff_reader_and_validate(reader, field_name)?,
|
||||
BucketAggregationType::DateHistogram(DateHistogramAggregationReq {
|
||||
field: field_name,
|
||||
..
|
||||
}) => get_ff_reader_and_validate(reader, field_name)?,
|
||||
}) => get_ff_reader_and_validate(reader, field_name, Cardinality::SingleValue)?,
|
||||
BucketAggregationType::Terms(TermsAggregation {
|
||||
field: field_name, ..
|
||||
}) => {
|
||||
str_dict_column = reader.fast_fields().str(field_name)?;
|
||||
get_ff_reader_and_validate(reader, field_name)?
|
||||
let field = reader.schema().get_field(field_name)?;
|
||||
inverted_index = Some(reader.inverted_index(field)?);
|
||||
get_ff_reader_and_validate(reader, field_name, Cardinality::MultiValues)?
|
||||
}
|
||||
};
|
||||
let sub_aggregation = sub_aggregation.clone();
|
||||
@@ -79,12 +106,15 @@ impl BucketAggregationWithAccessor {
|
||||
sub_aggregation: get_aggs_with_accessor_and_validate(
|
||||
&sub_aggregation,
|
||||
reader,
|
||||
&limits.clone(),
|
||||
bucket_count.clone(),
|
||||
max_bucket_count,
|
||||
)?,
|
||||
bucket_agg: bucket.clone(),
|
||||
str_dict_column,
|
||||
limits,
|
||||
column_block_accessor: Default::default(),
|
||||
inverted_index,
|
||||
bucket_count: BucketCount {
|
||||
bucket_count,
|
||||
max_bucket_count,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -93,9 +123,8 @@ impl BucketAggregationWithAccessor {
|
||||
#[derive(Clone)]
|
||||
pub struct MetricAggregationWithAccessor {
|
||||
pub metric: MetricAggregation,
|
||||
pub field_type: ColumnType,
|
||||
pub accessor: Column<u64>,
|
||||
pub column_block_accessor: ColumnBlockAccessor<u64>,
|
||||
pub field_type: Type,
|
||||
pub accessor: Arc<dyn Column>,
|
||||
}
|
||||
|
||||
impl MetricAggregationWithAccessor {
|
||||
@@ -110,13 +139,15 @@ impl MetricAggregationWithAccessor {
|
||||
| MetricAggregation::Min(MinAggregation { field: field_name })
|
||||
| MetricAggregation::Stats(StatsAggregation { field: field_name })
|
||||
| MetricAggregation::Sum(SumAggregation { field: field_name }) => {
|
||||
let (accessor, field_type) = get_ff_reader_and_validate(reader, field_name)?;
|
||||
let (accessor, field_type) =
|
||||
get_ff_reader_and_validate(reader, field_name, Cardinality::SingleValue)?;
|
||||
|
||||
Ok(MetricAggregationWithAccessor {
|
||||
accessor,
|
||||
accessor: accessor
|
||||
.into_single()
|
||||
.expect("unexpected fast field cardinality"),
|
||||
field_type,
|
||||
metric: metric.clone(),
|
||||
column_block_accessor: Default::default(),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -126,7 +157,8 @@ impl MetricAggregationWithAccessor {
|
||||
pub(crate) fn get_aggs_with_accessor_and_validate(
|
||||
aggs: &Aggregations,
|
||||
reader: &SegmentReader,
|
||||
limits: &AggregationLimits,
|
||||
bucket_count: Rc<AtomicU32>,
|
||||
max_bucket_count: u32,
|
||||
) -> crate::Result<AggregationsWithAccessor> {
|
||||
let mut metrics = vec![];
|
||||
let mut buckets = vec![];
|
||||
@@ -138,7 +170,8 @@ pub(crate) fn get_aggs_with_accessor_and_validate(
|
||||
&bucket.bucket_agg,
|
||||
&bucket.sub_aggregation,
|
||||
reader,
|
||||
limits.clone(),
|
||||
Rc::clone(&bucket_count),
|
||||
max_bucket_count,
|
||||
)?,
|
||||
)),
|
||||
Aggregation::Metric(metric) => metrics.push((
|
||||
@@ -157,15 +190,32 @@ pub(crate) fn get_aggs_with_accessor_and_validate(
|
||||
fn get_ff_reader_and_validate(
|
||||
reader: &SegmentReader,
|
||||
field_name: &str,
|
||||
) -> crate::Result<(columnar::Column<u64>, ColumnType)> {
|
||||
cardinality: Cardinality,
|
||||
) -> crate::Result<(FastFieldAccessor, Type)> {
|
||||
let field = reader.schema().get_field(field_name)?;
|
||||
let field_type = reader.schema().get_field_entry(field).field_type();
|
||||
|
||||
if let Some((_ff_type, field_cardinality)) = type_and_cardinality(field_type) {
|
||||
if cardinality != field_cardinality {
|
||||
return Err(TantivyError::InvalidArgument(format!(
|
||||
"Invalid field cardinality on field {} expected {:?}, but got {:?}",
|
||||
field_name, cardinality, field_cardinality
|
||||
)));
|
||||
}
|
||||
} else {
|
||||
return Err(TantivyError::InvalidArgument(format!(
|
||||
"Only fast fields of type f64, u64, i64 are supported, but got {:?} ",
|
||||
field_type.value_type()
|
||||
)));
|
||||
};
|
||||
|
||||
let ff_fields = reader.fast_fields();
|
||||
let ff_field_with_type = ff_fields
|
||||
.u64_lenient_with_type(field_name)?
|
||||
.unwrap_or_else(|| {
|
||||
(
|
||||
Column::build_empty_column(reader.num_docs()),
|
||||
ColumnType::U64,
|
||||
)
|
||||
});
|
||||
Ok(ff_field_with_type)
|
||||
match cardinality {
|
||||
Cardinality::SingleValue => ff_fields
|
||||
.u64_lenient(field_name)
|
||||
.map(|field| (FastFieldAccessor::Single(field), field_type.value_type())),
|
||||
Cardinality::MultiValues => ff_fields
|
||||
.u64s_lenient(field_name)
|
||||
.map(|field| (FastFieldAccessor::Multi(field), field_type.value_type())),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,8 +11,8 @@ use super::agg_req::BucketAggregationInternal;
|
||||
use super::bucket::GetDocCount;
|
||||
use super::intermediate_agg_result::{IntermediateBucketResult, IntermediateMetricResult};
|
||||
use super::metric::{SingleMetricResult, Stats};
|
||||
use super::segment_agg_result::AggregationLimits;
|
||||
use super::Key;
|
||||
use crate::schema::Schema;
|
||||
use crate::TantivyError;
|
||||
|
||||
#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
|
||||
@@ -20,13 +20,6 @@ use crate::TantivyError;
|
||||
pub struct AggregationResults(pub FxHashMap<String, AggregationResult>);
|
||||
|
||||
impl AggregationResults {
|
||||
pub(crate) fn get_bucket_count(&self) -> u64 {
|
||||
self.0
|
||||
.values()
|
||||
.map(|agg| agg.get_bucket_count())
|
||||
.sum::<u64>()
|
||||
}
|
||||
|
||||
pub(crate) fn get_value_from_aggregation(
|
||||
&self,
|
||||
name: &str,
|
||||
@@ -55,13 +48,6 @@ pub enum AggregationResult {
|
||||
}
|
||||
|
||||
impl AggregationResult {
|
||||
pub(crate) fn get_bucket_count(&self) -> u64 {
|
||||
match self {
|
||||
AggregationResult::BucketResult(bucket) => bucket.get_bucket_count(),
|
||||
AggregationResult::MetricResult(_) => 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn get_value_from_aggregation(
|
||||
&self,
|
||||
_name: &str,
|
||||
@@ -168,28 +154,12 @@ pub enum BucketResult {
|
||||
}
|
||||
|
||||
impl BucketResult {
|
||||
pub(crate) fn get_bucket_count(&self) -> u64 {
|
||||
match self {
|
||||
BucketResult::Range { buckets } => {
|
||||
buckets.iter().map(|bucket| bucket.get_bucket_count()).sum()
|
||||
}
|
||||
BucketResult::Histogram { buckets } => {
|
||||
buckets.iter().map(|bucket| bucket.get_bucket_count()).sum()
|
||||
}
|
||||
BucketResult::Terms {
|
||||
buckets,
|
||||
sum_other_doc_count: _,
|
||||
doc_count_error_upper_bound: _,
|
||||
} => buckets.iter().map(|bucket| bucket.get_bucket_count()).sum(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn empty_from_req(
|
||||
req: &BucketAggregationInternal,
|
||||
limits: &AggregationLimits,
|
||||
schema: &Schema,
|
||||
) -> crate::Result<Self> {
|
||||
let empty_bucket = IntermediateBucketResult::empty_from_req(&req.bucket_agg);
|
||||
empty_bucket.into_final_bucket_result(req, limits)
|
||||
empty_bucket.into_final_bucket_result(req, schema)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -204,15 +174,6 @@ pub enum BucketEntries<T> {
|
||||
HashMap(FxHashMap<String, T>),
|
||||
}
|
||||
|
||||
impl<T> BucketEntries<T> {
|
||||
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = &T> + 'a> {
|
||||
match self {
|
||||
BucketEntries::Vec(vec) => Box::new(vec.iter()),
|
||||
BucketEntries::HashMap(map) => Box::new(map.values()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// This is the default entry for a bucket, which contains a key, count, and optionally
|
||||
/// sub-aggregations.
|
||||
///
|
||||
@@ -252,11 +213,6 @@ pub struct BucketEntry {
|
||||
/// Sub-aggregations in this bucket.
|
||||
pub sub_aggregation: AggregationResults,
|
||||
}
|
||||
impl BucketEntry {
|
||||
pub(crate) fn get_bucket_count(&self) -> u64 {
|
||||
1 + self.sub_aggregation.get_bucket_count()
|
||||
}
|
||||
}
|
||||
impl GetDocCount for &BucketEntry {
|
||||
fn doc_count(&self) -> u64 {
|
||||
self.doc_count
|
||||
@@ -320,8 +276,3 @@ pub struct RangeBucketEntry {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub to_as_string: Option<String>,
|
||||
}
|
||||
impl RangeBucketEntry {
|
||||
pub(crate) fn get_bucket_count(&self) -> u64 {
|
||||
1 + self.sub_aggregation.get_bucket_count()
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user