reifydb_codec/frame/encoding/

dict.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

use std::{collections::HashMap, str};

use arrow_array::{Array, LargeBinaryArray, LargeStringArray};
use reifydb_value::value::blob::Blob;

use crate::{error::DecodeError, frame::format::dict_index_width_to_flags, tag::ValueKind};

pub struct DictEncoded {
	pub data: Vec<u8>,

	pub extra: Vec<u8>,

	pub type_code: u8,

	pub flags_bits: u8,
}

pub fn try_dict_encode_utf8(container: &LargeStringArray, min_ratio: f64) -> Option<DictEncoded> {
	let row_count = container.len();
	if row_count == 0 {
		return None;
	}

	let mut dict_map: HashMap<&str, u32> = HashMap::new();
	let mut dict_entries: Vec<&str> = Vec::new();

	for s in (0..row_count).map(|i| container.value(i)) {
		if !dict_map.contains_key(s) {
			let idx = dict_entries.len() as u32;
			dict_map.insert(s, idx);
			dict_entries.push(s);
		}
	}

	let dict_count = dict_entries.len();

	if dict_count as f64 >= row_count as f64 * min_ratio {
		return None;
	}

	let (index_width, flags_bits) = if dict_count <= 255 {
		(1usize, dict_index_width_to_flags(1))
	} else if dict_count <= 65535 {
		(2, dict_index_width_to_flags(2))
	} else {
		(4, dict_index_width_to_flags(4))
	};

	let mut data = Vec::with_capacity(row_count * index_width);
	for s in (0..row_count).map(|i| container.value(i)) {
		let idx = dict_map[s];
		match index_width {
			1 => data.push(idx as u8),
			2 => data.extend_from_slice(&(idx as u16).to_le_bytes()),
			4 => data.extend_from_slice(&idx.to_le_bytes()),
			_ => unreachable!(),
		}
	}

	let mut extra = Vec::new();
	extra.extend_from_slice(&(dict_count as u32).to_le_bytes());

	let mut offset: u32 = 0;
	extra.extend_from_slice(&offset.to_le_bytes());
	for entry in &dict_entries {
		offset += entry.len() as u32;
		extra.extend_from_slice(&offset.to_le_bytes());
	}

	for entry in &dict_entries {
		extra.extend_from_slice(entry.as_bytes());
	}

	Some(DictEncoded {
		data,
		extra,
		type_code: ValueKind::Utf8.byte(),
		flags_bits,
	})
}

pub fn try_dict_encode_blob(container: &LargeBinaryArray, min_ratio: f64) -> Option<DictEncoded> {
	let row_count = container.len();
	if row_count == 0 {
		return None;
	}

	let mut dict_map: HashMap<&[u8], u32> = HashMap::new();
	let mut dict_entries: Vec<&[u8]> = Vec::new();

	for bytes in (0..row_count).map(|i| container.value(i)) {
		if !dict_map.contains_key(bytes) {
			let idx = dict_entries.len() as u32;
			dict_map.insert(bytes, idx);
			dict_entries.push(bytes);
		}
	}

	let dict_count = dict_entries.len();

	if dict_count as f64 >= row_count as f64 * min_ratio {
		return None;
	}

	let (index_width, flags_bits) = if dict_count <= 255 {
		(1usize, dict_index_width_to_flags(1))
	} else if dict_count <= 65535 {
		(2, dict_index_width_to_flags(2))
	} else {
		(4, dict_index_width_to_flags(4))
	};

	let mut data = Vec::with_capacity(row_count * index_width);
	for bytes in (0..row_count).map(|i| container.value(i)) {
		let idx = dict_map[bytes];
		match index_width {
			1 => data.push(idx as u8),
			2 => data.extend_from_slice(&(idx as u16).to_le_bytes()),
			4 => data.extend_from_slice(&idx.to_le_bytes()),
			_ => unreachable!(),
		}
	}

	let mut extra = Vec::new();
	extra.extend_from_slice(&(dict_count as u32).to_le_bytes());

	let mut offset: u32 = 0;
	extra.extend_from_slice(&offset.to_le_bytes());
	for entry in &dict_entries {
		offset += entry.len() as u32;
		extra.extend_from_slice(&offset.to_le_bytes());
	}
	for entry in &dict_entries {
		extra.extend_from_slice(entry);
	}

	Some(DictEncoded {
		data,
		extra,
		type_code: ValueKind::Blob.byte(),
		flags_bits,
	})
}

pub fn decode_dict_utf8(
	data: &[u8],
	extra: &[u8],
	row_count: usize,
	index_width: usize,
) -> Result<Vec<String>, DecodeError> {
	let dict_entries = decode_dict_table(extra)?;

	let mut values = Vec::with_capacity(row_count);
	for i in 0..row_count {
		let idx = read_index(data, i, index_width) as usize;
		if idx >= dict_entries.len() {
			return Err(DecodeError::InvalidData(format!(
				"dict index {} out of range (dict has {} entries)",
				idx,
				dict_entries.len()
			)));
		}
		values.push(dict_entries[idx].clone());
	}

	Ok(values)
}

pub fn decode_dict_blob(
	data: &[u8],
	extra: &[u8],
	row_count: usize,
	index_width: usize,
) -> Result<Vec<Blob>, DecodeError> {
	let dict_entries = decode_dict_table_bytes(extra)?;

	let mut values = Vec::with_capacity(row_count);
	for i in 0..row_count {
		let idx = read_index(data, i, index_width) as usize;
		if idx >= dict_entries.len() {
			return Err(DecodeError::InvalidData(format!(
				"dict index {} out of range (dict has {} entries)",
				idx,
				dict_entries.len()
			)));
		}
		values.push(Blob::new(dict_entries[idx].clone()));
	}

	Ok(values)
}

fn decode_dict_table(extra: &[u8]) -> Result<Vec<String>, DecodeError> {
	if extra.len() < 4 {
		return Err(DecodeError::InvalidData("dict table too short".to_string()));
	}

	let dict_count = u32::from_le_bytes([extra[0], extra[1], extra[2], extra[3]]) as usize;
	let offsets_start = 4;
	let offsets_end = offsets_start + (dict_count + 1) * 4;

	if extra.len() < offsets_end {
		return Err(DecodeError::InvalidData("dict offsets truncated".to_string()));
	}

	let mut offsets = Vec::with_capacity(dict_count + 1);
	for i in 0..=dict_count {
		let pos = offsets_start + i * 4;
		offsets.push(u32::from_le_bytes([extra[pos], extra[pos + 1], extra[pos + 2], extra[pos + 3]]) as usize);
	}

	let data_start = offsets_end;
	let mut entries = Vec::with_capacity(dict_count);
	for i in 0..dict_count {
		let start = data_start + offsets[i];
		let end = data_start + offsets[i + 1];
		let s = str::from_utf8(&extra[start..end])
			.map_err(|e| DecodeError::InvalidData(format!("invalid dict UTF-8: {}", e)))?;
		entries.push(s.to_string());
	}

	Ok(entries)
}

pub fn decode_dict_table_bytes(extra: &[u8]) -> Result<Vec<Vec<u8>>, DecodeError> {
	if extra.len() < 4 {
		return Err(DecodeError::InvalidData("dict table too short".to_string()));
	}

	let dict_count = u32::from_le_bytes([extra[0], extra[1], extra[2], extra[3]]) as usize;
	let offsets_start = 4;
	let offsets_end = offsets_start + (dict_count + 1) * 4;

	if extra.len() < offsets_end {
		return Err(DecodeError::InvalidData("dict offsets truncated".to_string()));
	}

	let mut offsets = Vec::with_capacity(dict_count + 1);
	for i in 0..=dict_count {
		let pos = offsets_start + i * 4;
		offsets.push(u32::from_le_bytes([extra[pos], extra[pos + 1], extra[pos + 2], extra[pos + 3]]) as usize);
	}

	let data_start = offsets_end;
	let mut entries = Vec::with_capacity(dict_count);
	for i in 0..dict_count {
		let start = data_start + offsets[i];
		let end = data_start + offsets[i + 1];
		entries.push(extra[start..end].to_vec());
	}

	Ok(entries)
}

pub fn read_index(data: &[u8], i: usize, width: usize) -> u32 {
	let off = i * width;
	match width {
		1 => data[off] as u32,
		2 => u16::from_le_bytes([data[off], data[off + 1]]) as u32,
		4 => u32::from_le_bytes([data[off], data[off + 1], data[off + 2], data[off + 3]]),
		_ => unreachable!(),
	}
}