<commit_before>use crate::CodecId;
use crate::CodecReader;
use crate::FastFieldDataAccess;
use crate::FastFieldSerializerEstimate;
use crate::FastFieldStats;
use std::io::{self, Read, Write};
use std::ops::Sub;
use tantivy_bitpacker::compute_num_bits;
use tantivy_bitpacker::BitPacker;

use common::BinarySerializable;
use common::FixedSize;
use tantivy_bitpacker::BitUnpacker;

/// Depending on the field type, a different
/// fast field is required.
#[derive(Clone)]
pub struct LinearinterpolFastFieldReader {
    bit_unpacker: BitUnpacker,
    pub footer: LinearInterpolFooter,
    pub slope: f32,
}

#[derive(Clone, Debug)]
pub struct LinearInterpolFooter {
    pub relative_max_value: u64,
    pub offset: u64,
    pub first_val: u64,
    pub last_val: u64,
    pub num_vals: u64,
    pub min_value: u64,
    pub max_value: u64,
}

impl BinarySerializable for LinearInterpolFooter {
    fn serialize<W: Write>(&self, write: &mut W) -> io::Result<()> {
        self.relative_max_value.serialize(write)?;
        self.offset.serialize(write)?;
        self.first_val.serialize(write)?;
        self.last_val.serialize(write)?;
        self.num_vals.serialize(write)?;
        self.min_value.serialize(write)?;
        self.max_value.serialize(write)?;
        Ok(())
    }

    fn deserialize<R: Read>(reader: &mut R) -> io::Result<LinearInterpolFooter> {
        Ok(LinearInterpolFooter {
            relative_max_value: u64::deserialize(reader)?,
            offset: u64::deserialize(reader)?,
            first_val: u64::deserialize(reader)?,
            last_val: u64::deserialize(reader)?,
            num_vals: u64::deserialize(reader)?,
            min_value: u64::deserialize(reader)?,
            max_value: u64::deserialize(reader)?,
        })
    }
}

impl FixedSize for LinearInterpolFooter {
    const SIZE_IN_BYTES: usize = 56;
}

impl CodecReader for LinearinterpolFastFieldReader {
    /// Opens a fast field given a file.
    fn open_from_bytes(bytes: &[u8]) -> io::Result<Self> {
        let (_data, mut footer) = bytes.split_at(bytes.len() - LinearInterpolFooter::SIZE_IN_BYTES);
        let footer = LinearInterpolFooter::deserialize(&mut footer)?;
        let slope = get_slope(footer.first_val, footer.last_val, footer.num_vals);

        let num_bits = compute_num_bits(footer.relative_max_value);
        let bit_unpacker = BitUnpacker::new(num_bits);
        Ok(LinearinterpolFastFieldReader {
            bit_unpacker,
            footer,
            slope,
        })
    }
    #[inline]
    fn get_u64(&self, doc: u64, data: &[u8]) -> u64 {
        let calculated_value = get_calculated_value(self.footer.first_val, doc, self.slope);
        (calculated_value + self.bit_unpacker.get(doc, &data)) - self.footer.offset
    }

    #[inline]
    fn min_value(&self) -> u64 {
        self.footer.min_value
    }
    #[inline]
    fn max_value(&self) -> u64 {
        self.footer.max_value
    }
}

/// Fastfield serializer, which tries to guess values by linear interpolation
/// and stores the difference bitpacked.
pub struct LinearInterpolFastFieldSerializer {}

impl LinearInterpolFastFieldSerializer {
    /// Creates a new fast field serializer.
    pub fn create(
        write: &mut impl Write,
        fastfield_accessor: &impl FastFieldDataAccess,
        stats: FastFieldStats,
        data_iter: impl Iterator<Item = u64>,
        data_iter1: impl Iterator<Item = u64>,
    ) -> io::Result<()> {
        assert!(stats.min_value <= stats.max_value);

        let first_val = fastfield_accessor.get(0);
        let last_val = fastfield_accessor.get(stats.num_vals as u32 - 1);
        let slope = get_slope(first_val, last_val, stats.num_vals);
        // calculate offset to ensure all values are positive
        let mut offset = 0;
        let mut rel_positive_max = 0;
        for (pos, actual_value) in data_iter1.enumerate() {
            let calculated_value = get_calculated_value(first_val, pos as u64, slope);
            if calculated_value > actual_value {
                // negative value we need to apply an offset
                // we ignore negative values in the max value calculation, because negative values
                // will be offset to 0
                offset = offset.max(calculated_value - actual_value);
            } else {
                //positive value no offset reuqired
                rel_positive_max = rel_positive_max.max(actual_value - calculated_value);
            }
        }

        // rel_positive_max will be adjusted by offset
        let relative_max_value = rel_positive_max + offset;

        let num_bits = compute_num_bits(relative_max_value);
        let mut bit_packer = BitPacker::new();
        for (pos, val) in data_iter.enumerate() {
            let calculated_value = get_calculated_value(first_val, pos as u64, slope);
            let diff = (val + offset) - calculated_value;
            bit_packer.write(diff, num_bits, write)?;
        }
        bit_packer.close(write)?;

        let footer = LinearInterpolFooter {
            relative_max_value,
            offset,
            first_val,
            last_val,
            num_vals: stats.num_vals,
            min_value: stats.min_value,
            max_value: stats.max_value,
        };
        footer.serialize(write)?;
        Ok(())
    }
}
fn get_slope(first_val: u64, last_val: u64, num_vals: u64) -> f32 {
    (last_val as f32 - first_val as f32) / (num_vals as u64 - 1) as f32
}

fn get_calculated_value(first_val: u64, pos: u64, slope: f32) -> u64 {
    first_val + (pos as f32 * slope) as u64
}
impl FastFieldSerializerEstimate for LinearInterpolFastFieldSerializer {
    /// estimation for linear interpolation is hard because, you don't know
    /// where the local maxima are for the deviation of the calculated value and
    /// the offset is also unknown.
    fn estimate(fastfield_accessor: &impl FastFieldDataAccess, stats: FastFieldStats) -> f32 {
        if stats.max_value > i64::MAX as u64 / 2 || stats.num_vals < 3 {
            return f32::MAX; //disable compressor for this case
        }
        let first_val = fastfield_accessor.get(0);
        let last_val = fastfield_accessor.get(stats.num_vals as u32 - 1);
        let slope = get_slope(first_val, last_val, stats.num_vals);

        // let's sample at 10%, 20%, 25%, 50%, 75%, 90%
        let num_vals = stats.num_vals as f32 / 100.0;
        let sample_positions = [
            (num_vals * 10.0) as usize,
            (num_vals * 20.0) as usize,
            (num_vals * 25.0) as usize,
            (num_vals * 50.0) as usize,
            (num_vals * 75.0) as usize,
            (num_vals * 90.0) as usize,
        ];

        let max_distance = sample_positions
            .iter()
            .map(|pos| {
                let calculated_value = get_calculated_value(first_val, *pos as u64, slope);
                let actual_value = fastfield_accessor.get(*pos as u32);
                distance(calculated_value, actual_value)
            })
            .max()<commit_msg>use uniform distribution sampling<commit_after>use crate::CodecId;
use crate::CodecReader;
use crate::FastFieldDataAccess;
use crate::FastFieldSerializerEstimate;
use crate::FastFieldStats;
use std::io::{self, Read, Write};
use std::ops::Sub;
use tantivy_bitpacker::compute_num_bits;
use tantivy_bitpacker::BitPacker;

use common::BinarySerializable;
use common::FixedSize;
use tantivy_bitpacker::BitUnpacker;

/// Depending on the field type, a different
/// fast field is required.
#[derive(Clone)]
pub struct LinearinterpolFastFieldReader {
    bit_unpacker: BitUnpacker,
    pub footer: LinearInterpolFooter,
    pub slope: f32,
}

#[derive(Clone, Debug)]
pub struct LinearInterpolFooter {
    pub relative_max_value: u64,
    pub offset: u64,
    pub first_val: u64,
    pub last_val: u64,
    pub num_vals: u64,
    pub min_value: u64,
    pub max_value: u64,
}

impl BinarySerializable for LinearInterpolFooter {
    fn serialize<W: Write>(&self, write: &mut W) -> io::Result<()> {
        self.relative_max_value.serialize(write)?;
        self.offset.serialize(write)?;
        self.first_val.serialize(write)?;
        self.last_val.serialize(write)?;
        self.num_vals.serialize(write)?;
        self.min_value.serialize(write)?;
        self.max_value.serialize(write)?;
        Ok(())
    }

    fn deserialize<R: Read>(reader: &mut R) -> io::Result<LinearInterpolFooter> {
        Ok(LinearInterpolFooter {
            relative_max_value: u64::deserialize(reader)?,
            offset: u64::deserialize(reader)?,
            first_val: u64::deserialize(reader)?,
            last_val: u64::deserialize(reader)?,
            num_vals: u64::deserialize(reader)?,
            min_value: u64::deserialize(reader)?,
            max_value: u64::deserialize(reader)?,
        })
    }
}

impl FixedSize for LinearInterpolFooter {
    const SIZE_IN_BYTES: usize = 56;
}

impl CodecReader for LinearinterpolFastFieldReader {
    /// Opens a fast field given a file.
    fn open_from_bytes(bytes: &[u8]) -> io::Result<Self> {
        let (_data, mut footer) = bytes.split_at(bytes.len() - LinearInterpolFooter::SIZE_IN_BYTES);
        let footer = LinearInterpolFooter::deserialize(&mut footer)?;
        let slope = get_slope(footer.first_val, footer.last_val, footer.num_vals);

        let num_bits = compute_num_bits(footer.relative_max_value);
        let bit_unpacker = BitUnpacker::new(num_bits);
        Ok(LinearinterpolFastFieldReader {
            bit_unpacker,
            footer,
            slope,
        })
    }
    #[inline]
    fn get_u64(&self, doc: u64, data: &[u8]) -> u64 {
        let calculated_value = get_calculated_value(self.footer.first_val, doc, self.slope);
        (calculated_value + self.bit_unpacker.get(doc, &data)) - self.footer.offset
    }

    #[inline]
    fn min_value(&self) -> u64 {
        self.footer.min_value
    }
    #[inline]
    fn max_value(&self) -> u64 {
        self.footer.max_value
    }
}

/// Fastfield serializer, which tries to guess values by linear interpolation
/// and stores the difference bitpacked.
pub struct LinearInterpolFastFieldSerializer {}

impl LinearInterpolFastFieldSerializer {
    /// Creates a new fast field serializer.
    pub fn create(
        write: &mut impl Write,
        fastfield_accessor: &impl FastFieldDataAccess,
        stats: FastFieldStats,
        data_iter: impl Iterator<Item = u64>,
        data_iter1: impl Iterator<Item = u64>,
    ) -> io::Result<()> {
        assert!(stats.min_value <= stats.max_value);

        let first_val = fastfield_accessor.get(0);
        let last_val = fastfield_accessor.get(stats.num_vals as u32 - 1);
        let slope = get_slope(first_val, last_val, stats.num_vals);
        // calculate offset to ensure all values are positive
        let mut offset = 0;
        let mut rel_positive_max = 0;
        for (pos, actual_value) in data_iter1.enumerate() {
            let calculated_value = get_calculated_value(first_val, pos as u64, slope);
            if calculated_value > actual_value {
                // negative value we need to apply an offset
                // we ignore negative values in the max value calculation, because negative values
                // will be offset to 0
                offset = offset.max(calculated_value - actual_value);
            } else {
                //positive value no offset reuqired
                rel_positive_max = rel_positive_max.max(actual_value - calculated_value);
            }
        }

        // rel_positive_max will be adjusted by offset
        let relative_max_value = rel_positive_max + offset;

        let num_bits = compute_num_bits(relative_max_value);
        let mut bit_packer = BitPacker::new();
        for (pos, val) in data_iter.enumerate() {
            let calculated_value = get_calculated_value(first_val, pos as u64, slope);
            let diff = (val + offset) - calculated_value;
            bit_packer.write(diff, num_bits, write)?;
        }
        bit_packer.close(write)?;

        let footer = LinearInterpolFooter {
            relative_max_value,
            offset,
            first_val,
            last_val,
            num_vals: stats.num_vals,
            min_value: stats.min_value,
            max_value: stats.max_value,
        };
        footer.serialize(write)?;
        Ok(())
    }
}
fn get_slope(first_val: u64, last_val: u64, num_vals: u64) -> f32 {
    (last_val as f32 - first_val as f32) / (num_vals as u64 - 1) as f32
}

fn get_calculated_value(first_val: u64, pos: u64, slope: f32) -> u64 {
    first_val + (pos as f32 * slope) as u64
}
impl FastFieldSerializerEstimate for LinearInterpolFastFieldSerializer {
    /// estimation for linear interpolation is hard because, you don't know
    /// where the local maxima are for the deviation of the calculated value and
    /// the offset is also unknown.
    fn estimate(fastfield_accessor: &impl FastFieldDataAccess, stats: FastFieldStats) -> f32 {
        if stats.max_value > i64::MAX as u64 / 2 || stats.num_vals < 3 {
            return f32::MAX; //disable compressor for this case
        }
        let first_val = fastfield_accessor.get(0);
        let last_val = fastfield_accessor.get(stats.num_vals as u32 - 1);
        let slope = get_slope(first_val, last_val, stats.num_vals);

        // let's sample at 0%, 5%, 10% .. 95%, 100%
        let num_vals = stats.num_vals as f32 / 100.0;
        let sample_positions = (0..20)
            .map(|pos| (num_vals * pos as f32 * 5.0) as usize)
            .collect::<Vec<_>>();

        let max_distance = sample_positions
            .iter()
            .map(|pos| {
                let calculated_value = get_calculated_value(first_val, *pos as u64, slope);
                let actual_value = fastfield_accessor.get(*pos as u32);
                distance(calculated_value, actual_value)
            })
            .max()
{'fraction_non_alphanumeric': 0.10366642703091301, 'fraction_numerical': 0.0196980589503954, 'mean_word_length': 2.5623559539052496, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4109147', 'n_tokens_mistral': 4494, 'n_tokens_neox': 4282, 'n_words': 1063}