use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use crate::datatypes::values::Value;
use crate::graph::schema::TypeSchema;
use crate::graph::storage::interner::{InternedKey, StringInterner, STRIP_PROPERTIES};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct ColumnarRow {
row_id: u32,
}
impl ColumnarRow {
#[inline]
pub(crate) fn new(row_id: u32) -> Self {
ColumnarRow { row_id }
}
#[inline]
pub(crate) fn row_id(&self) -> u32 {
self.row_id
}
}
pub(crate) enum PropertyStorage {
Map(HashMap<InternedKey, Value>),
Compact {
schema: Arc<TypeSchema>,
values: Vec<Value>,
},
Columnar(ColumnarRow),
}
pub(crate) enum PropertyKeyIter<'a> {
Map {
inner: std::collections::hash_map::Keys<'a, InternedKey, Value>,
interner: &'a StringInterner,
},
Compact {
slots: &'a [InternedKey],
values: &'a [Value],
slot_idx: usize,
interner: &'a StringInterner,
},
Columnar(std::vec::IntoIter<&'a str>),
}
impl<'a> Iterator for PropertyKeyIter<'a> {
type Item = &'a str;
#[inline]
fn next(&mut self) -> Option<&'a str> {
match self {
PropertyKeyIter::Map { inner, interner } => inner.next().map(|k| interner.resolve(*k)),
PropertyKeyIter::Compact {
slots,
values,
slot_idx,
interner,
} => loop {
let i = *slot_idx;
if i >= slots.len() {
return None;
}
*slot_idx = i + 1;
if values.get(i).is_some_and(|v| !matches!(v, Value::Null)) {
return Some(interner.resolve(slots[i]));
}
},
PropertyKeyIter::Columnar(iter) => iter.next(),
}
}
}
impl PropertyStorage {
#[inline]
pub(in crate::graph::storage) fn get(&self, key: InternedKey) -> Option<Cow<'_, Value>> {
match self {
PropertyStorage::Map(map) => map.get(&key).map(Cow::Borrowed),
PropertyStorage::Compact { schema, values } => schema
.slot(key)
.and_then(|slot| values.get(slot as usize))
.filter(|v| !matches!(v, Value::Null))
.map(Cow::Borrowed),
PropertyStorage::Columnar(_) => None,
}
}
#[inline]
pub(in crate::graph::storage) fn get_value(&self, key: InternedKey) -> Option<Value> {
match self {
PropertyStorage::Map(map) => map.get(&key).cloned(),
PropertyStorage::Compact { schema, values } => schema
.slot(key)
.and_then(|slot| values.get(slot as usize))
.filter(|v| !matches!(v, Value::Null))
.cloned(),
PropertyStorage::Columnar(_) => None,
}
}
#[inline]
pub(in crate::graph::storage) fn contains(&self, key: InternedKey) -> bool {
self.get(key).is_some()
}
#[inline]
pub(in crate::graph::storage) fn str_prop_eq(
&self,
key: InternedKey,
target: &str,
) -> Option<bool> {
match self {
PropertyStorage::Map(map) => map
.get(&key)
.map(|v| matches!(v, Value::String(s) if s == target)),
PropertyStorage::Compact { schema, values } => schema
.slot(key)
.and_then(|slot| values.get(slot as usize))
.filter(|v| !matches!(v, Value::Null))
.map(|v| matches!(v, Value::String(s) if s == target)),
PropertyStorage::Columnar(_) => None,
}
}
pub(in crate::graph::storage) fn insert(&mut self, key: InternedKey, value: Value) {
match self {
PropertyStorage::Map(map) => {
map.insert(key, value);
}
PropertyStorage::Compact { schema, values } => {
let slot = if let Some(s) = schema.slot(key) {
s as usize
} else {
let s = Arc::make_mut(schema).add_key(key) as usize;
s
};
if slot >= values.len() {
values.resize(slot + 1, Value::Null);
}
values[slot] = value;
}
PropertyStorage::Columnar(_) => debug_assert!(
false,
"columnar property write must go through GraphWrite::set_node_property"
),
}
}
pub(in crate::graph::storage) fn insert_if_absent(&mut self, key: InternedKey, value: Value) {
match self {
PropertyStorage::Map(map) => {
map.entry(key).or_insert(value);
}
PropertyStorage::Compact { schema, values } => {
if let Some(slot) = schema.slot(key) {
let slot = slot as usize;
if slot < values.len() {
if matches!(values[slot], Value::Null) {
values[slot] = value;
}
} else {
values.resize(slot + 1, Value::Null);
values[slot] = value;
}
} else {
let slot = Arc::make_mut(schema).add_key(key) as usize;
if slot >= values.len() {
values.resize(slot + 1, Value::Null);
}
values[slot] = value;
}
}
PropertyStorage::Columnar(_) => debug_assert!(
false,
"columnar property write must go through GraphWrite::set_node_property_if_absent"
),
}
}
pub(in crate::graph::storage) fn remove(&mut self, key: InternedKey) -> Option<Value> {
match self {
PropertyStorage::Map(map) => map.remove(&key),
PropertyStorage::Compact { schema, values } => schema.slot(key).and_then(|slot| {
let slot = slot as usize;
if slot < values.len() {
let old = std::mem::replace(&mut values[slot], Value::Null);
if matches!(old, Value::Null) {
None
} else {
Some(old)
}
} else {
None
}
}),
PropertyStorage::Columnar(_) => {
debug_assert!(
false,
"columnar property removal must go through GraphWrite::remove_node_property"
);
None
}
}
}
pub(in crate::graph::storage) fn replace_all(
&mut self,
pairs: impl IntoIterator<Item = (InternedKey, Value)>,
) {
match self {
PropertyStorage::Map(map) => {
map.clear();
map.extend(pairs);
}
PropertyStorage::Compact { schema, values } => {
for v in values.iter_mut() {
*v = Value::Null;
}
for (key, value) in pairs {
let slot = if let Some(s) = schema.slot(key) {
s as usize
} else {
Arc::make_mut(schema).add_key(key) as usize
};
if slot >= values.len() {
values.resize(slot + 1, Value::Null);
}
values[slot] = value;
}
}
PropertyStorage::Columnar(_) => debug_assert!(
false,
"columnar property replace must go through GraphWrite::replace_node_properties"
),
}
}
pub(in crate::graph::storage) fn len(&self) -> usize {
match self {
PropertyStorage::Map(map) => map.len(),
PropertyStorage::Compact { values, .. } => {
values.iter().filter(|v| !matches!(v, Value::Null)).count()
}
PropertyStorage::Columnar(_) => 0,
}
}
pub fn drain_to_interned_pairs(
&mut self,
_interner: &StringInterner,
) -> Vec<(InternedKey, Value)> {
match std::mem::replace(self, PropertyStorage::Map(HashMap::new())) {
PropertyStorage::Map(map) => map.into_iter().collect(),
PropertyStorage::Compact { schema, values } => schema
.slots
.iter()
.zip(values)
.filter(|(_, v)| !matches!(v, Value::Null))
.map(|(ik, v)| (*ik, v))
.collect(),
PropertyStorage::Columnar { .. } => {
Vec::new()
}
}
}
pub(in crate::graph::storage) fn keys<'a>(
&'a self,
interner: &'a StringInterner,
) -> PropertyKeyIter<'a> {
match self {
PropertyStorage::Map(map) => PropertyKeyIter::Map {
inner: map.keys(),
interner,
},
PropertyStorage::Compact { schema, values } => PropertyKeyIter::Compact {
slots: &schema.slots,
values,
slot_idx: 0,
interner,
},
PropertyStorage::Columnar(_) => PropertyKeyIter::Columnar(Vec::new().into_iter()),
}
}
pub fn from_compact(
pairs: impl IntoIterator<Item = (InternedKey, Value)>,
schema: &Arc<TypeSchema>,
) -> Self {
let mut values = vec![Value::Null; schema.len()];
for (key, value) in pairs {
if let Some(slot) = schema.slot(key) {
values[slot as usize] = value;
}
}
PropertyStorage::Compact {
schema: Arc::clone(schema),
values,
}
}
}
impl Clone for PropertyStorage {
fn clone(&self) -> Self {
match self {
PropertyStorage::Map(map) => PropertyStorage::Map(map.clone()),
PropertyStorage::Compact { schema, values } => PropertyStorage::Compact {
schema: Arc::clone(schema),
values: values.clone(),
},
PropertyStorage::Columnar(row) => PropertyStorage::Columnar(*row),
}
}
}
impl std::fmt::Debug for PropertyStorage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PropertyStorage::Map(map) => f.debug_tuple("Map").field(map).finish(),
PropertyStorage::Compact { values, .. } => {
f.debug_tuple("Compact").field(values).finish()
}
PropertyStorage::Columnar(row) => f
.debug_struct("Columnar")
.field("row_id", &row.row_id())
.finish(),
}
}
}
impl PartialEq for PropertyStorage {
fn eq(&self, other: &Self) -> bool {
fn collect_entries(ps: &PropertyStorage) -> Vec<(InternedKey, Value)> {
match ps {
PropertyStorage::Map(map) => {
let mut entries: Vec<_> = map.iter().map(|(&k, v)| (k, v.clone())).collect();
entries.sort_by_key(|(k, _)| k.as_u64());
entries
}
PropertyStorage::Compact { schema, values } => {
let mut entries: Vec<_> = schema
.slots
.iter()
.enumerate()
.filter_map(|(i, &ik)| {
values.get(i).and_then(|v| {
if matches!(v, Value::Null) {
None
} else {
Some((ik, v.clone()))
}
})
})
.collect();
entries.sort_by_key(|(k, _)| k.as_u64());
entries
}
PropertyStorage::Columnar(_) => Vec::new(),
}
}
collect_entries(self) == collect_entries(other)
}
}
impl Serialize for PropertyStorage {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
use serde::ser::SerializeMap;
if STRIP_PROPERTIES.with(|cell| cell.get()) {
return serializer.serialize_map(Some(0))?.end();
}
match self {
PropertyStorage::Map(map) => map.serialize(serializer),
PropertyStorage::Compact { schema, values } => {
let count = values.iter().filter(|v| !matches!(v, Value::Null)).count();
let mut map_ser = serializer.serialize_map(Some(count))?;
for (i, ik) in schema.slots.iter().enumerate() {
if let Some(v) = values.get(i) {
if !matches!(v, Value::Null) {
map_ser.serialize_entry(ik, v)?;
}
}
}
map_ser.end()
}
PropertyStorage::Columnar(_) => {
debug_assert!(
STRIP_PROPERTIES.with(|cell| cell.get()),
"serializing a columnar node without STRIP_PROPERTIES would write \
an empty property map: the save path must persist the type's \
ColumnStore separately (see io/file.rs) or convert first"
);
serializer.serialize_map(Some(0))?.end()
}
}
}
}
impl<'de> Deserialize<'de> for PropertyStorage {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let map = HashMap::<InternedKey, Value>::deserialize(deserializer)?;
Ok(PropertyStorage::Map(map))
}
}
impl PropertyStorage {
#[inline]
pub(crate) fn columnar_row_id(&self) -> Option<u32> {
match self {
PropertyStorage::Columnar(row) => Some(row.row_id()),
_ => None,
}
}
}