use arrow_array::{Array, PrimitiveArray, types::ArrowPrimitiveType};
use oxgraph_topology::{
ElementIndex, ElementWeight, IncidenceBase, IncidenceIndex, IncidenceWeight, RelationIndex,
RelationWeight, TopologyBase,
};
use crate::{
model::{IdFamily, PropertyError, PropertyLayer, PropertyLayerData, ensure_no_nulls},
width::{AxisIndex, ElementAxis, IncidenceAxis, PropertyAxis, PropertyIndex, RelationAxis},
};
#[derive(Clone, Copy, Debug)]
pub struct GraphPropertyLayers<'view, Id, NodeIndex, EdgeIndex>
where
NodeIndex: PropertyIndex,
EdgeIndex: PropertyIndex,
{
pub element: &'view [PropertyLayer<Id, NodeIndex>],
pub relation: &'view [PropertyLayer<Id, EdgeIndex>],
}
#[derive(Clone, Copy, Debug)]
pub struct HyperPropertyLayers<'view, Id, VertexIndex, RelationIndex, IncidenceIndex>
where
VertexIndex: PropertyIndex,
RelationIndex: PropertyIndex,
IncidenceIndex: PropertyIndex,
{
pub element: &'view [PropertyLayer<Id, VertexIndex>],
pub relation: &'view [PropertyLayer<Id, RelationIndex>],
pub incidence: &'view [PropertyLayer<Id, IncidenceIndex>],
}
macro_rules! impl_axis_topology_base {
($storage:ident, $axis:ty, $index_trait:ident) => {
impl<T, Id, I, P> TopologyBase for $storage<'_, $axis, T, Id, I, P>
where
T: $index_trait,
I: PropertyIndex,
P: ArrowPrimitiveType,
{
type ElementId = T::ElementId;
type RelationId = T::RelationId;
}
};
}
macro_rules! impl_dense_axis_weight {
(
$axis:ty, $index_trait:ident, $id_ty:ident, $index_fn:ident,
$weight_trait:ident, $weight_fn:ident
) => {
impl<T, Id, I, P> $weight_trait for DenseWeights<'_, $axis, T, Id, I, P>
where
T: $index_trait,
I: PropertyIndex,
P: ArrowPrimitiveType,
P::Native: Copy,
{
type Weight = P::Native;
fn $weight_fn(&self, id: Self::$id_ty) -> Self::Weight {
self.values.value(self.topology.$index_fn(id))
}
}
};
}
macro_rules! impl_sparse_axis_weight {
(
$axis:ty, $index_trait:ident, $id_ty:ident, $index_fn:ident,
$weight_trait:ident, $weight_fn:ident
) => {
impl<T, Id, I, P> $weight_trait for SparseWeights<'_, $axis, T, Id, I, P>
where
T: $index_trait,
I: PropertyIndex,
P: ArrowPrimitiveType,
P::Native: Copy,
{
type Weight = P::Native;
fn $weight_fn(&self, id: Self::$id_ty) -> Self::Weight {
sparse_value::<I, P>(
self.indices,
self.values,
self.default,
self.topology.$index_fn(id),
)
}
}
};
}
macro_rules! impl_axis_incidence_base {
($storage:ident) => {
impl<T, Id, I, P> IncidenceBase for $storage<'_, IncidenceAxis, T, Id, I, P>
where
T: IncidenceIndex,
I: PropertyIndex,
P: ArrowPrimitiveType,
{
type IncidenceId = T::IncidenceId;
type Role = T::Role;
}
};
}
pub struct DenseWeights<'view, A, T, Id, I, P>
where
A: PropertyAxis,
I: PropertyIndex,
P: ArrowPrimitiveType,
{
topology: &'view T,
values: &'view PrimitiveArray<P>,
property: core::marker::PhantomData<(A, Id, I)>,
}
impl<'view, A, T, Id, I, P> DenseWeights<'view, A, T, Id, I, P>
where
A: PropertyAxis,
T: AxisIndex<A>,
I: PropertyIndex,
P: ArrowPrimitiveType,
{
pub fn new(
topology: &'view T,
layer: &'view PropertyLayer<Id, I>,
) -> Result<Self, PropertyError> {
let values = validate_dense_primitive_selection::<Id, I, P>(
layer,
A::id_family(),
topology.axis_bound(),
)?;
Ok(Self {
topology,
values,
property: core::marker::PhantomData,
})
}
}
impl_axis_topology_base!(DenseWeights, ElementAxis, ElementIndex);
impl_axis_topology_base!(DenseWeights, RelationAxis, RelationIndex);
impl_axis_topology_base!(DenseWeights, IncidenceAxis, IncidenceIndex);
impl_axis_incidence_base!(DenseWeights);
impl_dense_axis_weight!(
ElementAxis,
ElementIndex,
ElementId,
element_index,
ElementWeight,
element_weight
);
impl_dense_axis_weight!(
RelationAxis,
RelationIndex,
RelationId,
relation_index,
RelationWeight,
relation_weight
);
impl_dense_axis_weight!(
IncidenceAxis,
IncidenceIndex,
IncidenceId,
incidence_index,
IncidenceWeight,
incidence_weight
);
pub struct SparseWeights<'view, A, T, Id, I, P>
where
A: PropertyAxis,
I: PropertyIndex,
P: ArrowPrimitiveType,
{
topology: &'view T,
indices: &'view PrimitiveArray<I::ArrowType>,
values: &'view PrimitiveArray<P>,
default: P::Native,
property: core::marker::PhantomData<(A, Id)>,
}
impl<'view, A, T, Id, I, P> SparseWeights<'view, A, T, Id, I, P>
where
A: PropertyAxis,
T: AxisIndex<A>,
I: PropertyIndex,
P: ArrowPrimitiveType,
P::Native: Copy,
{
pub fn new(
topology: &'view T,
layer: &'view PropertyLayer<Id, I>,
) -> Result<Self, PropertyError> {
let (indices, values, default) = validate_sparse_primitive_selection::<I, P, Id>(
layer,
A::id_family(),
topology.axis_bound(),
)?;
Ok(Self {
topology,
indices,
values,
default,
property: core::marker::PhantomData,
})
}
}
impl_axis_topology_base!(SparseWeights, ElementAxis, ElementIndex);
impl_axis_topology_base!(SparseWeights, RelationAxis, RelationIndex);
impl_axis_topology_base!(SparseWeights, IncidenceAxis, IncidenceIndex);
impl_axis_incidence_base!(SparseWeights);
impl_sparse_axis_weight!(
ElementAxis,
ElementIndex,
ElementId,
element_index,
ElementWeight,
element_weight
);
impl_sparse_axis_weight!(
RelationAxis,
RelationIndex,
RelationId,
relation_index,
RelationWeight,
relation_weight
);
impl_sparse_axis_weight!(
IncidenceAxis,
IncidenceIndex,
IncidenceId,
incidence_index,
IncidenceWeight,
incidence_weight
);
fn validate_dense_primitive_selection<Id, I, P>(
layer: &PropertyLayer<Id, I>,
expected: IdFamily,
required: usize,
) -> Result<&PrimitiveArray<P>, PropertyError>
where
I: PropertyIndex,
P: ArrowPrimitiveType,
{
if layer.descriptor().id_family != expected {
return Err(PropertyError::IdFamilyMismatch {
expected,
actual: layer.descriptor().id_family,
});
}
if layer.len() < required {
return Err(PropertyError::LayerTooShort {
required,
actual: layer.len(),
});
}
let PropertyLayerData::Dense { values } = layer.data() else {
return Err(PropertyError::ExpectedDenseStorage {
name: layer.descriptor().name.clone(),
});
};
let primitive = values
.as_any()
.downcast_ref::<PrimitiveArray<P>>()
.ok_or_else(|| PropertyError::ArrowTypeMismatch {
name: layer.descriptor().name.clone(),
})?;
ensure_no_nulls(primitive)?;
Ok(primitive)
}
type SparsePrimitiveSelection<'layer, I, P> = (
&'layer PrimitiveArray<<I as PropertyIndex>::ArrowType>,
&'layer PrimitiveArray<P>,
<P as ArrowPrimitiveType>::Native,
);
fn validate_sparse_primitive_selection<I, P, Id>(
layer: &PropertyLayer<Id, I>,
expected: IdFamily,
required: usize,
) -> Result<SparsePrimitiveSelection<'_, I, P>, PropertyError>
where
I: PropertyIndex,
P: ArrowPrimitiveType,
P::Native: Copy,
{
if layer.descriptor().id_family != expected {
return Err(PropertyError::IdFamilyMismatch {
expected,
actual: layer.descriptor().id_family,
});
}
if layer.len() < required {
return Err(PropertyError::LayerTooShort {
required,
actual: layer.len(),
});
}
let PropertyLayerData::Sparse {
indices,
values,
default,
} = layer.data()
else {
return Err(PropertyError::ExpectedSparseStorage {
name: layer.descriptor().name.clone(),
});
};
let Some(default_array) = default else {
return Err(PropertyError::SparseNullMissingNotTotal {
name: layer.descriptor().name.clone(),
});
};
let primitive = values
.as_any()
.downcast_ref::<PrimitiveArray<P>>()
.ok_or_else(|| PropertyError::ArrowTypeMismatch {
name: layer.descriptor().name.clone(),
})?;
ensure_no_nulls(primitive)?;
let default_primitive = default_array
.as_any()
.downcast_ref::<PrimitiveArray<P>>()
.ok_or_else(|| PropertyError::ArrowTypeMismatch {
name: layer.descriptor().name.clone(),
})?;
if default_primitive.len() != 1 || default_primitive.is_null(0) {
return Err(PropertyError::DefaultPolicyMismatch {
name: layer.descriptor().name.clone(),
});
}
Ok((indices.as_ref(), primitive, default_primitive.value(0)))
}
fn sparse_value<I, P>(
indices: &PrimitiveArray<I::ArrowType>,
values: &PrimitiveArray<P>,
default: P::Native,
index: usize,
) -> P::Native
where
I: PropertyIndex,
P: ArrowPrimitiveType,
P::Native: Copy,
{
let Some(target) = I::from_usize(index) else {
return default;
};
let mut low = 0_usize;
let mut high = indices.len();
while low < high {
let mid = low + ((high - low) / 2);
let value = indices.value(mid);
if value < target {
low = mid + 1;
} else {
high = mid;
}
}
if low < indices.len() && indices.value(low) == target {
values.value(low)
} else {
default
}
}