use std::sync::Arc;
use vortex_array::IntoArray;
use vortex_array::aggregate_fn::AggregateFnVTableExt;
use vortex_array::aggregate_fn::EmptyOptions;
use vortex_array::arrays::ExtensionArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::StructArray;
use vortex_array::dtype::DType;
use vortex_array::dtype::FieldNames;
use vortex_array::dtype::extension::ExtDType;
use vortex_array::validity::Validity;
use vortex_error::VortexResult;
use vortex_layout::layouts::zoned::zone_map::ZoneMap;
use crate::aggregate_fn::GeometryAabb;
use crate::extension::GeoMetadata;
use crate::extension::Rect;
pub(super) fn aabb_zone_map(point_dtype: &DType, boxes: &[[f64; 4]]) -> VortexResult<ZoneMap> {
let aabb_fn = GeometryAabb.bind(EmptyOptions);
let col = |i: usize| PrimitiveArray::from_iter(boxes.iter().map(move |b| b[i])).into_array();
let storage = StructArray::try_new(
["xmin", "ymin", "xmax", "ymax"].into(),
vec![col(0), col(1), col(2), col(3)],
boxes.len(),
Validity::AllValid,
)?
.into_array();
let box_dtype =
ExtDType::<Rect>::try_new(GeoMetadata::default(), storage.dtype().clone())?.erased();
let aabbs = ExtensionArray::try_new(box_dtype, storage)?.into_array();
let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?;
ZoneMap::try_new(
point_dtype.clone(),
zone_array,
Arc::new([aabb_fn]),
1,
boxes.len() as u64,
)
}
pub(super) fn empty_zone_map(point_dtype: &DType) -> VortexResult<ZoneMap> {
ZoneMap::try_new(
point_dtype.clone(),
StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable)?,
Arc::new([]),
1,
2,
)
}