Don't pass strings to the choose_array_encoder method when choosing a…

…n encoder for dict indices
lancedb · broccoliSpicy · Sep 27, 2024 · Sep 16, 2024 · Sep 16, 2024 · Sep 17, 2024
commit 403e89df1efd76f08225270a194612b4fbf5686a
diff --git a/rust/lance-encoding/src/encoder.rs b/rust/lance-encoding/src/encoder.rs
@@ -3,7 +3,7 @@
 use std::{collections::HashMap, env, sync::Arc};
 
 use arrow::array::AsArray;
-use arrow_array::{Array, ArrayRef, RecordBatch};
+use arrow_array::{Array, ArrayRef, RecordBatch, UInt8Array};
 use arrow_schema::DataType;
 use bytes::{Bytes, BytesMut};
 use futures::future::BoxFuture;
@@ -269,7 +269,11 @@ impl CoreArrayEncodingStrategy {
             DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
                 if use_dict_encoding {
                     let dict_indices_encoder = Self::choose_array_encoder(
-                        arrays,
+                        // We need to pass arrays to this method to figure out what kind of compression to
+                        // use but we haven't actually calculated the indices yet.  For now, we just assume
+                        // worst case and use the full range.  In the future maybe we can pass in statistics
+                        // instead of the actual data
+                        &[Arc::new(UInt8Array::from_iter_values(0_u8..255_u8))],
                         &DataType::UInt8,
                         data_size,
                         false,

diff --git a/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs b/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs

 macro_rules! encode_fixed_width {
    ($self:expr, $unpacked:expr, $data_type:ty, $buffer_index:expr) => {{
        let num_chunks = ($unpacked.num_values + ELEMS_PER_CHUNK - 1) / ELEMS_PER_CHUNK;
        let num_full_chunks = $unpacked.num_values / ELEMS_PER_CHUNK;
        let uncompressed_bit_width = std::mem::size_of::<$data_type>() as u64 * 8;
        
        let DataBlock::FixedWidth(mut unpacked) = data else {
            return Err(Error::InvalidInput {
                source: "Bitpacking only supports fixed width data blocks".into(),
                location: location!(),
            });
        };
    
    bytes_idx_to_range_indices: &[Vec<std::ops::Range<u64>>],
    num_rows: u64,
 ) -> LanceBuffer {
    match uncompressed_bits_per_value {
        8 => {
            let mut decompressed: Vec<u8> = Vec::with_capacity(num_rows as usize);
            let packed_chunk_size: usize = ELEMS_PER_CHUNK as usize * compressed_bit_width as usize / 8;
                            compressed_bit_width as usize,
                            chunk,
                            &mut decompress_chunk_buf[..ELEMS_PER_CHUNK as usize],
                        );
                    }
                    loop {
                        if curr_range_start + ELEMS_PER_CHUNK < bytes_idx_to_range_indices[i][ranges_idx].end {
                            curr_range_start += this_part_len;
                            break;
                        } else {
                            let this_part_len =
                                bytes_idx_to_range_indices[i][ranges_idx].end - curr_range_start;
                            decompressed.extend_from_slice(
                                &decompress_chunk_buf[(curr_range_start % ELEMS_PER_CHUNK) as usize..]
            }
            LanceBuffer::Owned(decompressed)
        }

        16 => {
            let mut decompressed: Vec<u16> = Vec::with_capacity(num_rows as usize);
            let packed_chunk_size_in_byte: usize = (ELEMS_PER_CHUNK * compressed_bit_width) as usize / 8;
                            compressed_bit_width as usize,
                            chunk,
                            &mut decompress_chunk_buf,
                        );
                    }
                    loop {
                        if curr_range_start + ELEMS_PER_CHUNK < bytes_idx_to_range_indices[i][ranges_idx].end {
                            // we know this chunk has only data of this range
                            break;
                        } else {
                            let this_part_len =
                                bytes_idx_to_range_indices[i][ranges_idx].end - curr_range_start;
                            decompressed.extend_from_slice(
                                &decompress_chunk_buf[(curr_range_start % ELEMS_PER_CHUNK) as usize..]
            }
            LanceBuffer::reinterpret_vec(decompressed).to_owned()
        }

        32 => {
            let mut decompressed: Vec<u32> = Vec::with_capacity(num_rows as usize);
            let packed_chunk_size_in_byte: usize = (ELEMS_PER_CHUNK * compressed_bit_width) as usize / 8;
                            compressed_bit_width as usize,
                            chunk,
                            &mut decompress_chunk_buf,
                        );
                    }
                    loop {
                        if curr_range_start + ELEMS_PER_CHUNK < bytes_idx_to_range_indices[i][ranges_idx].end {

	macro_rules! encode_fixed_width {
	($self:expr, $unpacked:expr, $data_type:ty, $buffer_index:expr) => {{
	let num_chunks = ($unpacked.num_values + ELEMS_PER_CHUNK - 1) / ELEMS_PER_CHUNK;
Check warning on line 199 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	let num_full_chunks = $unpacked.num_values / ELEMS_PER_CHUNK;
	let uncompressed_bit_width = std::mem::size_of::<$data_type>() as u64 * 8;

	let DataBlock::FixedWidth(mut unpacked) = data else {
	return Err(Error::InvalidInput {
	source: "Bitpacking only supports fixed width data blocks".into(),
	location: location!(),
Check warning on line 290 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	});
	};

	bytes_idx_to_range_indices: &[Vec<std::ops::Range<u64>>],
	num_rows: u64,
	) -> LanceBuffer {
	match uncompressed_bits_per_value {
Check warning on line 444 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	8 => {
	let mut decompressed: Vec<u8> = Vec::with_capacity(num_rows as usize);
	let packed_chunk_size: usize = ELEMS_PER_CHUNK as usize * compressed_bit_width as usize / 8;
	compressed_bit_width as usize,
	chunk,
	&mut decompress_chunk_buf[..ELEMS_PER_CHUNK as usize],
	);
Check warning on line 460 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	}
	loop {
	if curr_range_start + ELEMS_PER_CHUNK < bytes_idx_to_range_indices[i][ranges_idx].end {
	curr_range_start += this_part_len;
	break;
	} else {
	let this_part_len =
Check warning on line 471 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	bytes_idx_to_range_indices[i][ranges_idx].end - curr_range_start;
	decompressed.extend_from_slice(
	&decompress_chunk_buf[(curr_range_start % ELEMS_PER_CHUNK) as usize..]
	}
	LanceBuffer::Owned(decompressed)
	}

Check warning on line 489 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	16 => {
	let mut decompressed: Vec<u16> = Vec::with_capacity(num_rows as usize);
	let packed_chunk_size_in_byte: usize = (ELEMS_PER_CHUNK * compressed_bit_width) as usize / 8;
	compressed_bit_width as usize,
	chunk,
	&mut decompress_chunk_buf,
	);
Check warning on line 507 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	}
	loop {
	if curr_range_start + ELEMS_PER_CHUNK < bytes_idx_to_range_indices[i][ranges_idx].end {
	// we know this chunk has only data of this range
	break;
	} else {
	let this_part_len =
Check warning on line 521 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	bytes_idx_to_range_indices[i][ranges_idx].end - curr_range_start;
	decompressed.extend_from_slice(
	&decompress_chunk_buf[(curr_range_start % ELEMS_PER_CHUNK) as usize..]
	}
	LanceBuffer::reinterpret_vec(decompressed).to_owned()
	}

Check warning on line 539 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	32 => {
	let mut decompressed: Vec<u32> = Vec::with_capacity(num_rows as usize);
	let packed_chunk_size_in_byte: usize = (ELEMS_PER_CHUNK * compressed_bit_width) as usize / 8;
	compressed_bit_width as usize,
	chunk,
	&mut decompress_chunk_buf,
	);
Check warning on line 557 in rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs GitHub Actions / linux-arm `Diff in /runner/_work/lance/lance/rust/lance-encoding/src/encodings/physical/bitpack_fastlanes.rs`
	}
	loop {
	if curr_range_start + ELEMS_PER_CHUNK < bytes_idx_to_range_indices[i][ranges_idx].end {