use std::io;
use super::{FileMetadata, PortableColumnSection};
const NODE_ROW_BYTES: u64 = 64;
const STRING_CELL_BYTES: u64 = 40;
const NUMERIC_CELL_BYTES: u64 = 16;
const BOOL_CELL_BYTES: u64 = 8;
const UNIQUE_ENTRY_BYTES: u64 = 48;
const POSTING_ENTRY_BYTES: u64 = 8;
const KEY_DENSITY_DIVISOR: u64 = 8;
const COMPOSITE_KEY_VEC_BYTES: u64 = 24;
const STRING_VALUE_BYTES: u64 = 48;
const SCALAR_VALUE_BYTES: u64 = 16;
const DECOMPRESS_RATIO_NUMERATOR: u64 = 9;
const DECOMPRESS_RATIO_DENOMINATOR: u64 = 2;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LoadMemoryEstimate {
pub index_rebuild_bytes: u64,
pub section_heap_bytes: u64,
pub transient_peak_bytes: u64,
pub node_rows: u64,
pub declared_indexes: u32,
}
impl LoadMemoryEstimate {
pub fn total_settled_bytes(&self) -> u64 {
self.section_heap_bytes
.saturating_add(self.index_rebuild_bytes)
}
pub fn total_peak_bytes(&self) -> u64 {
self.total_settled_bytes()
.saturating_add(self.transient_peak_bytes)
}
pub fn projected_peak_bytes(&self, defer_index_rebuild: bool) -> u64 {
if defer_index_rebuild {
self.section_heap_bytes
.saturating_add(self.transient_peak_bytes)
} else {
self.total_peak_bytes()
}
}
}
fn cell_bytes(type_tag: &str) -> u64 {
match type_tag {
"bool" => BOOL_CELL_BYTES,
"int64" | "float64" | "date" | "uniqueid" => NUMERIC_CELL_BYTES,
_ => STRING_CELL_BYTES,
}
}
fn value_bytes(type_tag: &str) -> u64 {
match type_tag {
"int64" | "float64" | "date" | "uniqueid" | "bool" => SCALAR_VALUE_BYTES,
_ => STRING_VALUE_BYTES,
}
}
fn section_for<'a>(
metadata: &'a FileMetadata,
node_type: &str,
) -> Option<&'a PortableColumnSection> {
metadata
.column_sections
.iter()
.find(|section| section.type_name == node_type)
}
fn tag_for<'a>(metadata: &'a FileMetadata, node_type: &str, property: &str) -> Option<&'a str> {
section_for(metadata, node_type)
.and_then(|section| section.columns.get(property))
.map(String::as_str)
}
fn rows_for(metadata: &FileMetadata, node_type: &str) -> u64 {
section_for(metadata, node_type).map_or(0, |section| u64::from(section.row_count))
}
fn equality_index_row_bytes(
metadata: &FileMetadata,
node_type: &str,
properties: &[String],
) -> u64 {
let mut key_bytes: u64 = if properties.len() > 1 {
COMPOSITE_KEY_VEC_BYTES
} else {
0
};
for property in properties {
key_bytes += tag_for(metadata, node_type, property).map_or(STRING_VALUE_BYTES, value_bytes);
}
POSTING_ENTRY_BYTES + key_bytes / KEY_DENSITY_DIVISOR
}
fn unique_index_row_bytes(metadata: &FileMetadata, node_type: &str, properties: &[String]) -> u64 {
let mut bytes = UNIQUE_ENTRY_BYTES + COMPOSITE_KEY_VEC_BYTES;
for property in properties {
bytes += tag_for(metadata, node_type, property).map_or(STRING_VALUE_BYTES, value_bytes);
}
bytes
}
pub(super) fn estimate_from_metadata(metadata: &FileMetadata) -> LoadMemoryEstimate {
let mut node_rows: u64 = 0;
let mut section_heap_bytes: u64 = 0;
let mut compressed_total: u64 = metadata.topology_compressed_size;
let mut largest_section: u64 = metadata.topology_compressed_size;
for section in &metadata.column_sections {
let rows = u64::from(section.row_count);
node_rows += rows;
section_heap_bytes += rows * NODE_ROW_BYTES;
for type_tag in section.columns.values() {
section_heap_bytes += rows * cell_bytes(type_tag);
}
compressed_total += section.compressed_size;
largest_section = largest_section.max(section.compressed_size);
}
for optional in [
metadata.embeddings_compressed_size,
metadata.timeseries_compressed_size,
metadata.secondary_labels_compressed_size,
metadata.vector_index_compressed_size,
metadata.text_index_compressed_size,
] {
compressed_total += optional;
largest_section = largest_section.max(optional);
}
section_heap_bytes = section_heap_bytes.max(compressed_total);
let mut index_rebuild_bytes: u64 = 0;
let mut declared_indexes: u32 = 0;
for (node_type, property) in metadata
.property_index_keys
.iter()
.chain(metadata.range_index_keys.iter())
{
declared_indexes += 1;
index_rebuild_bytes += rows_for(metadata, node_type)
* equality_index_row_bytes(metadata, node_type, std::slice::from_ref(property));
}
for (node_type, properties) in &metadata.composite_index_keys {
declared_indexes += 1;
index_rebuild_bytes += rows_for(metadata, node_type)
* equality_index_row_bytes(metadata, node_type, properties);
}
for (node_type, properties) in &metadata.unique_constraint_keys {
declared_indexes += 1;
index_rebuild_bytes +=
rows_for(metadata, node_type) * unique_index_row_bytes(metadata, node_type, properties);
}
LoadMemoryEstimate {
index_rebuild_bytes,
section_heap_bytes,
transient_peak_bytes: largest_section * DECOMPRESS_RATIO_NUMERATOR
/ DECOMPRESS_RATIO_DENOMINATOR,
node_rows,
declared_indexes,
}
}
pub fn estimate_load_memory(path: &str) -> io::Result<LoadMemoryEstimate> {
if std::path::Path::new(path).is_dir() {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"'{path}' is a disk-mode graph directory, not a portable .kgl. Load memory is \
estimated from a .kgl's metadata head; a disk graph keeps its columns and \
indexes on disk and never rebuilds them at load, so the terms this reports \
would not describe it."
),
));
}
let head = super::read_metadata_head_from_file(path)?;
Ok(estimate_from_metadata(&head))
}
pub fn estimate_load_memory_bytes(data: &[u8]) -> io::Result<LoadMemoryEstimate> {
let head = super::read_metadata_head(data, "the byte buffer")?;
Ok(estimate_from_metadata(&head))
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
struct Fixture {
rows: u32,
columns: &'static [(&'static str, &'static str)],
topology_compressed: u64,
column_compressed: u64,
measured_settled_bytes: u64,
}
const ITEM_COLUMNS: &[(&str, &str)] = &[
("category", "string"),
("count", "int64"),
("region", "string"),
("score", "float64"),
("sku", "string"),
];
const INDEXED_500K: Fixture = Fixture {
rows: 500_000,
columns: ITEM_COLUMNS,
topology_compressed: 2_659_663,
column_compressed: 9_436_825,
measured_settled_bytes: 150_700_000,
};
const NOINDEX_500K: Fixture = Fixture {
rows: 500_000,
columns: ITEM_COLUMNS,
topology_compressed: 2_659_663,
column_compressed: 9_436_825,
measured_settled_bytes: 86_600_000,
};
fn metadata_for(fixture: &Fixture, indexed: bool) -> FileMetadata {
let mut metadata = FileMetadata {
topology_compressed_size: fixture.topology_compressed,
..Default::default()
};
metadata.column_sections.push(PortableColumnSection {
type_name: "Item".to_string(),
compressed_size: fixture.column_compressed,
row_count: fixture.rows,
columns: fixture
.columns
.iter()
.map(|(name, tag)| (name.to_string(), tag.to_string()))
.collect::<HashMap<String, String>>(),
});
if indexed {
metadata.property_index_keys = vec![
("Item".to_string(), "category".to_string()),
("Item".to_string(), "region".to_string()),
];
metadata.composite_index_keys = vec![(
"Item".to_string(),
vec!["category".to_string(), "region".to_string()],
)];
metadata.unique_constraint_keys = vec![("Item".to_string(), vec!["sku".to_string()])];
}
metadata
}
#[test]
fn estimate_is_within_a_factor_of_two_of_measured_footprint() {
for (name, fixture, indexed) in [
("indexed_500k", &INDEXED_500K, true),
("noindex_500k", &NOINDEX_500K, false),
] {
let estimate = estimate_from_metadata(&metadata_for(fixture, indexed));
let settled = estimate.total_settled_bytes() as f64;
let measured = fixture.measured_settled_bytes as f64;
let ratio = settled / measured;
assert!(
(0.5..=2.0).contains(&ratio),
"{name}: estimated {settled:.0} B against {measured:.0} B measured (×{ratio:.2})"
);
}
}
#[test]
fn index_term_matches_the_measured_index_term() {
let indexed = estimate_from_metadata(&metadata_for(&INDEXED_500K, true));
let noindex = estimate_from_metadata(&metadata_for(&NOINDEX_500K, false));
assert_eq!(noindex.index_rebuild_bytes, 0, "no declaration, no term");
assert_eq!(noindex.declared_indexes, 0);
assert_eq!(indexed.declared_indexes, 4);
assert_eq!(indexed.section_heap_bytes, noindex.section_heap_bytes);
let term = indexed.index_rebuild_bytes as f64;
assert!(
(55e6..=95e6).contains(&term),
"index term {term:.0} B is outside the 64.1-79.1 MB measured band's neighbourhood"
);
}
#[test]
fn deferring_indexes_would_remove_the_index_term() {
let indexed = estimate_from_metadata(&metadata_for(&INDEXED_500K, true));
assert_eq!(
indexed.total_settled_bytes() - indexed.index_rebuild_bytes,
indexed.section_heap_bytes
);
}
#[test]
fn an_index_on_a_type_with_no_section_contributes_nothing() {
let mut metadata = metadata_for(&NOINDEX_500K, false);
metadata.property_index_keys = vec![("Ghost".to_string(), "name".to_string())];
let estimate = estimate_from_metadata(&metadata);
assert_eq!(estimate.index_rebuild_bytes, 0);
assert_eq!(estimate.declared_indexes, 1);
}
#[test]
fn transient_is_the_largest_single_section() {
let estimate = estimate_from_metadata(&metadata_for(&NOINDEX_500K, false));
let expected = NOINDEX_500K.column_compressed * DECOMPRESS_RATIO_NUMERATOR
/ DECOMPRESS_RATIO_DENOMINATOR;
assert_eq!(estimate.transient_peak_bytes, expected);
assert_eq!(
estimate.total_peak_bytes(),
estimate.total_settled_bytes() + expected
);
}
#[test]
fn an_empty_graph_falls_back_to_the_compressed_floor() {
let metadata = FileMetadata {
topology_compressed_size: 4096,
..Default::default()
};
let estimate = estimate_from_metadata(&metadata);
assert_eq!(estimate.node_rows, 0);
assert_eq!(estimate.section_heap_bytes, 4096);
}
#[test]
fn an_unknown_column_tag_costs_a_string() {
assert_eq!(cell_bytes("from-a-newer-writer"), STRING_CELL_BYTES);
assert_eq!(value_bytes("from-a-newer-writer"), STRING_VALUE_BYTES);
}
}