diff --git a/vortex-geo/benches/envelope.rs b/vortex-geo/benches/envelope.rs index 6ae5d6404c2..ca356c2a585 100644 --- a/vortex-geo/benches/envelope.rs +++ b/vortex-geo/benches/envelope.rs @@ -25,6 +25,8 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::scalar_fn::EmptyOptions; use vortex_geo::scalar_fn::envelope::GeoEnvelope; use vortex_geo::test_harness::MultiPolygonRings; use vortex_geo::test_harness::geo_session; @@ -58,9 +60,9 @@ fn coin(i: usize) -> bool { /// Execute the envelope of `column` to completion. fn envelope(column: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { - GeoEnvelope::try_new_array(column.clone()) + GeoEnvelope + .try_new_array(column.len(), EmptyOptions, [column.clone()]) .unwrap() - .into_array() .execute::(ctx) .unwrap() .into_array() diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index 855a6af1867..e197fad7ebe 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -6,7 +6,6 @@ use geo::Contains; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; @@ -17,7 +16,6 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -31,17 +29,6 @@ use crate::scalar_fn::execute::execute_null_propagating; #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct GeoContains; -impl GeoContains { - /// A lazy `ScalarFnArray` computing per-row whether operand `a` contains operand `b`; - /// either may be constant. The output length is taken from `a`. - pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { - ScalarFnArray::try_new( - TypedScalarFnInstance::new(GeoContains, EmptyOptions).erased(), - vec![a, b], - ) - } -} - impl ScalarFnVTable for GeoContains { type Options = EmptyOptions; @@ -127,6 +114,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -177,7 +165,7 @@ mod tests { ) -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); - let contains = GeoContains::try_new_array(a, b)?.into_array(); + let contains = GeoContains.try_new_array(a.len(), EmptyOptions, [a, b])?; assert_arrays_eq!(contains, BoolArray::from_iter(expected), &mut ctx); Ok(()) } @@ -298,7 +286,8 @@ mod tests { let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; let points = nullable_point_column(vec![Some((2.0, 2.0)), None, Some((10.0, 10.0))])?; - let contains = GeoContains::try_new_array(container, points)?.into_array(); + let contains = + GeoContains.try_new_array(container.len(), EmptyOptions, [container, points])?; let expected = BoolArray::new( BitBuffer::from_iter([true, false, false]), @@ -318,7 +307,8 @@ mod tests { let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable(); let null_const = ConstantArray::new(Scalar::null(point_dtype), 2).into_array(); let points = point_column(vec![2.0, 10.0], vec![2.0, 10.0])?; - let contains = GeoContains::try_new_array(null_const, points)?.into_array(); + let contains = + GeoContains.try_new_array(null_const.len(), EmptyOptions, [null_const, points])?; let expected = BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); @@ -346,7 +336,8 @@ mod tests { None, Some((4.0, 4.0)), ])?; - let contains = GeoContains::try_new_array(container, contained)?.into_array(); + let contains = + GeoContains.try_new_array(container.len(), EmptyOptions, [container, contained])?; let expected = BoolArray::new( BitBuffer::from_iter([true, false, false, false]), @@ -365,7 +356,8 @@ mod tests { let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; let points = nullable_point_column(vec![None, None])?; - let contains = GeoContains::try_new_array(container, points)?.into_array(); + let contains = + GeoContains.try_new_array(container.len(), EmptyOptions, [container, points])?; let expected = BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); @@ -382,7 +374,8 @@ mod tests { let container = nullable_point_column(vec![Some((1.0, 1.0)), None])?; let contained = nullable_point_column(vec![None, Some((2.0, 2.0))])?; - let contains = GeoContains::try_new_array(container, contained)?.into_array(); + let contains = + GeoContains.try_new_array(container.len(), EmptyOptions, [container, contained])?; let expected = BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index dd209a712b4..28bf480e98e 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -7,7 +7,6 @@ use geo::Distance; use geo::Euclidean; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -19,7 +18,6 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -32,17 +30,6 @@ use crate::scalar_fn::execute::execute_null_propagating; #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct GeoDistance; -impl GeoDistance { - /// A lazy `ScalarFnArray` computing the per-row distance between operands `a` and `b`; either may - /// be constant. The output length is taken from `a`. - pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { - ScalarFnArray::try_new( - TypedScalarFnInstance::new(GeoDistance, EmptyOptions).erased(), - vec![a, b], - ) - } -} - impl ScalarFnVTable for GeoDistance { type Options = EmptyOptions; @@ -115,6 +102,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -159,7 +147,7 @@ mod tests { let a = point_column(vec![0.0, 3.0, 0.0, 3.0], vec![0.0, 0.0, 4.0, 4.0])?; let b = point_constant(0.0, 0.0, 4, &mut ctx)?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; assert_eq!(distances(distance, &mut ctx)?, vec![0.0, 3.0, 4.0, 5.0]); Ok(()) @@ -173,7 +161,7 @@ mod tests { let a = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?; let b = point_column(vec![3.0, 1.0], vec![4.0, 1.0])?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 0.0]); Ok(()) @@ -190,7 +178,7 @@ mod tests { let single = polygon_column(vec![vec![ring]])?.execute_scalar(0, &mut ctx)?; let square = ConstantArray::new(single, 2).into_array(); let points = point_column(vec![7.0, 2.0], vec![2.0, 2.0])?; - let distance = GeoDistance::try_new_array(points, square)?.into_array(); + let distance = GeoDistance.try_new_array(points.len(), EmptyOptions, [points, square])?; assert_eq!(distances(distance, &mut ctx)?, vec![3.0, 0.0]); Ok(()) @@ -204,7 +192,7 @@ mod tests { let a = point_constant(0.0, 0.0, 4, &mut ctx)?; let b = point_column(vec![0.0, 3.0, 0.0, 3.0], vec![0.0, 0.0, 4.0, 4.0])?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; assert_eq!(distances(distance, &mut ctx)?, vec![0.0, 3.0, 4.0, 5.0]); Ok(()) @@ -218,7 +206,7 @@ mod tests { let a = point_constant(0.0, 0.0, 3, &mut ctx)?; let b = point_constant(3.0, 4.0, 3, &mut ctx)?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 5.0, 5.0]); Ok(()) @@ -245,7 +233,7 @@ mod tests { let a = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((3.0, 4.0))])?; let b = point_constant(0.0, 0.0, 3, &mut ctx)?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; let expected = PrimitiveArray::new( vec![0.0f64, 0.0, 5.0], @@ -264,7 +252,7 @@ mod tests { let a = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((0.0, 0.0))])?; let b = nullable_point_column(vec![Some((3.0, 4.0)), Some((1.0, 1.0)), None])?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; let expected = PrimitiveArray::new( vec![5.0f64, 0.0, 0.0], @@ -284,7 +272,8 @@ mod tests { let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable(); let null_const = ConstantArray::new(Scalar::null(point_dtype), 3).into_array(); let b = point_column(vec![0.0, 3.0, 0.0], vec![0.0, 0.0, 4.0])?; - let distance = GeoDistance::try_new_array(null_const, b)?.into_array(); + let distance = + GeoDistance.try_new_array(null_const.len(), EmptyOptions, [null_const, b])?; let expected = PrimitiveArray::new(vec![0.0f64; 3], Validity::AllInvalid).into_array(); assert_arrays_eq!(distance, expected, &mut ctx); @@ -299,7 +288,7 @@ mod tests { let a = nullable_point_column(vec![None, None])?; let b = point_constant(0.0, 0.0, 2, &mut ctx)?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; let expected = PrimitiveArray::new(vec![0.0f64; 2], Validity::AllInvalid).into_array(); assert_arrays_eq!(distance, expected, &mut ctx); @@ -315,7 +304,7 @@ mod tests { let a = nullable_point_column(vec![Some((0.0, 0.0)), None])?; let b = nullable_point_column(vec![None, Some((1.0, 1.0))])?; - let distance = GeoDistance::try_new_array(a, b)?.into_array(); + let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?; let expected = PrimitiveArray::new(vec![0.0f64; 2], Validity::AllInvalid).into_array(); assert_arrays_eq!(distance, expected, &mut ctx); @@ -332,8 +321,8 @@ mod tests { let a = point_column(vec![], vec![])?; let b = point_column(vec![], vec![])?; - let result = GeoDistance::try_new_array(a, b)? - .into_array() + let result = GeoDistance + .try_new_array(a.len(), EmptyOptions, [a, b])? .execute::(&mut ctx)? .into_array(); diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs index 1859bde666c..35ad80f8a39 100644 --- a/vortex-geo/src/scalar_fn/envelope.rs +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -10,7 +10,6 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::struct_::StructArrayExt; @@ -25,7 +24,6 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; @@ -53,17 +51,6 @@ use crate::extension::validate_geometry_operands; #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct GeoEnvelope; -impl GeoEnvelope { - /// A lazy `ScalarFnArray` computing the per-row bounding box of geometry operand `a`, which may - /// be constant. The output length is taken from `a`. - pub fn try_new_array(a: ArrayRef) -> VortexResult { - ScalarFnArray::try_new( - TypedScalarFnInstance::new(GeoEnvelope, EmptyOptions).erased(), - vec![a], - ) - } -} - /// The output dtype: a nullable native 2-D box ([`Rect`], `geoarrow.box`) column. Nullable /// because rows without a box — null or empty geometries — are null. Metadata is defaulted. fn output_box_dtype() -> VortexResult> { @@ -231,6 +218,7 @@ mod tests { use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -259,7 +247,7 @@ mod tests { /// Execute a `GeoEnvelope` over `array`, returning the lazy box column. fn boxes(array: ArrayRef) -> VortexResult { - Ok(GeoEnvelope::try_new_array(array)?.into_array()) + GeoEnvelope.try_new_array(array.len(), EmptyOptions, [array]) } /// A point's box is degenerate: both corners are the point itself. diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index 3f3842e4e1c..81fba3a97e4 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -6,7 +6,6 @@ use geo::Intersects; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; @@ -17,7 +16,6 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -30,17 +28,6 @@ use crate::scalar_fn::execute::execute_null_propagating; #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct GeoIntersects; -impl GeoIntersects { - /// A lazy `ScalarFnArray` computing per-row whether operands `a` and `b` intersect; either may - /// be constant. The output length is taken from `a`. - pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { - ScalarFnArray::try_new( - TypedScalarFnInstance::new(GeoIntersects, EmptyOptions).erased(), - vec![a, b], - ) - } -} - impl ScalarFnVTable for GeoIntersects { type Options = EmptyOptions; @@ -126,6 +113,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -197,7 +185,7 @@ mod tests { ) -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); - let intersects = GeoIntersects::try_new_array(a, b)?.into_array(); + let intersects = GeoIntersects.try_new_array(a.len(), EmptyOptions, [a, b])?; assert_arrays_eq!(intersects, BoolArray::from_iter(expected), &mut ctx); Ok(()) } @@ -283,9 +271,11 @@ mod tests { let constant = geometry_constant(&donut(), 4)?; let column = materialize(constant.clone(), &mut ctx)?; + let probes = donut_probes()?; let against_constant = - GeoIntersects::try_new_array(donut_probes()?, constant)?.into_array(); - let pairwise = GeoIntersects::try_new_array(donut_probes()?, column)?.into_array(); + GeoIntersects.try_new_array(probes.len(), EmptyOptions, [probes, constant])?; + let probes = donut_probes()?; + let pairwise = GeoIntersects.try_new_array(probes.len(), EmptyOptions, [probes, column])?; assert_arrays_eq!(against_constant, pairwise, &mut ctx); Ok(()) @@ -321,7 +311,8 @@ mod tests { let points = nullable_point_column(vec![Some((2.0, 2.0)), None, Some((20.0, 20.0))])?; let query = geometry_constant(&donut(), 3)?; - let intersects = GeoIntersects::try_new_array(points, query)?.into_array(); + let intersects = + GeoIntersects.try_new_array(points.len(), EmptyOptions, [points, query])?; let expected = BoolArray::new( BitBuffer::from_iter([true, false, false]), @@ -341,7 +332,8 @@ mod tests { let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable(); let null_const = ConstantArray::new(Scalar::null(point_dtype), 2).into_array(); let points = point_column(vec![2.0, 20.0], vec![2.0, 20.0])?; - let intersects = GeoIntersects::try_new_array(null_const, points)?.into_array(); + let intersects = + GeoIntersects.try_new_array(null_const.len(), EmptyOptions, [null_const, points])?; let expected = BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); @@ -368,7 +360,7 @@ mod tests { None, Some((9.0, 9.0)), ])?; - let intersects = GeoIntersects::try_new_array(a, b)?.into_array(); + let intersects = GeoIntersects.try_new_array(a.len(), EmptyOptions, [a, b])?; let expected = BoolArray::new( BitBuffer::from_iter([true, false, false, false]), @@ -387,7 +379,8 @@ mod tests { let points = nullable_point_column(vec![None, None])?; let query = geometry_constant(&donut(), 2)?; - let intersects = GeoIntersects::try_new_array(points, query)?.into_array(); + let intersects = + GeoIntersects.try_new_array(points.len(), EmptyOptions, [points, query])?; let expected = BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); @@ -404,7 +397,7 @@ mod tests { let a = nullable_point_column(vec![Some((0.0, 0.0)), None])?; let b = nullable_point_column(vec![None, Some((1.0, 1.0))])?; - let intersects = GeoIntersects::try_new_array(a, b)?.into_array(); + let intersects = GeoIntersects.try_new_array(a.len(), EmptyOptions, [a, b])?; let expected = BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); diff --git a/vortex-geo/src/tests/rect.rs b/vortex-geo/src/tests/rect.rs index 865892493e7..fc8740a1737 100644 --- a/vortex-geo/src/tests/rect.rs +++ b/vortex-geo/src/tests/rect.rs @@ -13,8 +13,10 @@ use geoarrow::datatypes::Crs; use geoarrow::datatypes::Dimension as GeoArrowDimension; use geoarrow::datatypes::Metadata; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::scalar_fn::EmptyOptions; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -116,7 +118,6 @@ fn roundtrips_through_arrow() -> VortexResult<()> { #[test] fn scalar_functions_run_on_rect() -> VortexResult<()> { use vortex_array::Canonical; - use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; @@ -130,17 +131,19 @@ fn scalar_functions_run_on_rect() -> VortexResult<()> { let points = point_column(vec![5.0, 20.0], vec![5.0, 20.0])?; // Distance: 0 to the interior point, >0 to the exterior point. - let distance = GeoDistance::try_new_array(bbox.clone(), points.clone())?.into_array(); + let distance = + GeoDistance.try_new_array(bbox.len(), EmptyOptions, [bbox.clone(), points.clone()])?; let distance = distance.execute::(&mut ctx)?.into_primitive(); let distances = distance.as_slice::(); assert_eq!(distances[0], 0.0); assert!(distances[1] > 0.0); // Intersects / Contains: true for the interior point, false for the exterior one. - let intersects = GeoIntersects::try_new_array(bbox.clone(), points.clone())?.into_array(); + let intersects = + GeoIntersects.try_new_array(bbox.len(), EmptyOptions, [bbox.clone(), points.clone()])?; assert_arrays_eq!(intersects, BoolArray::from_iter([true, false]), &mut ctx); - let contains = GeoContains::try_new_array(bbox, points)?.into_array(); + let contains = GeoContains.try_new_array(bbox.len(), EmptyOptions, [bbox, points])?; assert_arrays_eq!(contains, BoolArray::from_iter([true, false]), &mut ctx); Ok(()) } diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index de6f0614471..ceb8b49698f 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -8,8 +8,8 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; @@ -23,7 +23,6 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::serde::ArrayChildren; use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -58,24 +57,6 @@ use crate::utils::validate_binary_tensor_float_inputs; #[derive(Clone)] pub struct CosineSimilarity; -impl CosineSimilarity { - /// Creates a new [`TypedScalarFnInstance`] wrapping the cosine similarity operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(CosineSimilarity, EmptyOptions) - } - - /// Constructs a [`ScalarFnArray`] that lazily computes the cosine similarity between `lhs` and - /// `rhs`. - /// - /// # Errors - /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). - pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(CosineSimilarity::new().erased(), vec![lhs, rhs]) - } -} - impl ScalarFnVTable for CosineSimilarity { type Options = EmptyOptions; @@ -141,9 +122,9 @@ impl ScalarFnVTable for CosineSimilarity { let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm::try_new_array(lhs_ref.clone())?; - let norm_rhs_arr = L2Norm::try_new_array(rhs_ref.clone())?; - let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; + let norm_lhs_arr = L2Norm.try_new_array(len, EmptyOptions, [lhs_ref.clone()])?; + let norm_rhs_arr = L2Norm.try_new_array(len, EmptyOptions, [rhs_ref.clone()])?; + let dot_arr = InnerProduct.try_new_array(len, EmptyOptions, [lhs_ref, rhs_ref])?; // Execute to get the inner product and norms of the arrays. We only fully decompress // because we need to perform special logic (guard against 0) during division. @@ -236,8 +217,8 @@ impl CosineSimilarity { // `L2Denorm` makes the normalized children authoritative, so their dot product is the // cosine similarity even for lossy storage wrappers, except that a zero stored norm still // represents a zero vector. - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized_l, normalized_r])? .execute(ctx)?; let norms_l: PrimitiveArray = norms_l.execute(ctx)?; let norms_r: PrimitiveArray = norms_r.execute(ctx)?; @@ -276,13 +257,14 @@ impl CosineSimilarity { let (normalized, denorm_norms) = extract_l2_denorm_children(denorm_ref); - let dot_arr = InnerProduct::try_new_array(normalized, plain_ref.clone())?; - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; + let dot_arr = + InnerProduct.try_new_array(len, EmptyOptions, [normalized, plain_ref.clone()])?; + let dot: PrimitiveArray = dot_arr.execute(ctx)?; let denorm_norms: PrimitiveArray = denorm_norms.execute(ctx)?; - let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; - let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; + let norm_arr = L2Norm.try_new_array(len, EmptyOptions, [plain_ref.clone()])?; + let plain_norm: PrimitiveArray = norm_arr.execute(ctx)?; // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. // TODO(connor): This can be written in a more SIMD-friendly manner. @@ -316,8 +298,9 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; + use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_error::VortexResult; @@ -333,8 +316,7 @@ mod tests { /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let result = CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs])?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; Ok(prim.as_slice::().to_vec()) @@ -498,8 +480,7 @@ mod tests { let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let result = CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs])?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; @@ -582,8 +563,7 @@ mod tests { let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); let rhs = L2Denorm::try_new_array(normalized_r, norms_r, &mut ctx)?.into_array(); - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let result = CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs])?; let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; assert!(prim.is_valid(0, &mut ctx)?); @@ -747,7 +727,8 @@ mod tests { #[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] #[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = CosineSimilarity::try_new_array(lhs.clone(), rhs.clone())?.into_array(); + let original = + CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); let metadata = plugin diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index f46b03ac94e..c12d8595eb7 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -9,9 +9,9 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; @@ -26,7 +26,6 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::serde::ArrayChildren; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -55,24 +54,6 @@ use crate::utils::validate_binary_tensor_float_inputs; #[derive(Clone)] pub struct InnerProduct; -impl InnerProduct { - /// Creates a new [`TypedScalarFnInstance`] wrapping the inner product operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(InnerProduct, EmptyOptions) - } - - /// Constructs a [`ScalarFnArray`] that lazily computes the inner product between `lhs` and - /// `rhs`. - /// - /// # Errors - /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). - pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(InnerProduct::new().erased(), vec![lhs, rhs]) - } -} - impl ScalarFnVTable for InnerProduct { type Options = EmptyOptions; @@ -219,8 +200,8 @@ impl InnerProduct { let norms_l: PrimitiveArray = norms_l.execute(ctx)?; let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized_l, normalized_r])? .execute(ctx)?; match_each_float_ptype!(dot.ptype(), |T| { @@ -249,8 +230,8 @@ impl InnerProduct { let (normalized, norms) = extract_l2_denorm_children(denorm_ref); let denorm_norms: PrimitiveArray = norms.execute(ctx)?; - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized, plain_ref.clone())? - .into_array() + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized, plain_ref.clone()])? .execute(ctx)?; match_each_float_ptype!(dot.ptype(), |T| { @@ -284,8 +265,9 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; + use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_error::VortexResult; @@ -299,8 +281,7 @@ mod tests { /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs])?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; Ok(prim.as_slice::().to_vec()) @@ -377,8 +358,7 @@ mod tests { let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs])?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; @@ -395,7 +375,7 @@ mod tests { fn rejects_non_extension_dtype() { let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); - let result = InnerProduct::try_new_array(lhs, rhs); + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); assert!(result.is_err()); } @@ -403,7 +383,7 @@ mod tests { fn rejects_mismatched_dtypes() -> VortexResult<()> { let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; let rhs = vector_array(2, &[3.0_f64, 4.0])?; - let result = InnerProduct::try_new_array(lhs, rhs); + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); assert!(result.is_err()); Ok(()) } @@ -470,8 +450,7 @@ mod tests { let lhs = L2Denorm::try_new_array(normalized_l, norms_l, &mut ctx)?.into_array(); let rhs = l2_denorm_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs])?; let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. @@ -485,7 +464,8 @@ mod tests { #[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] #[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = InnerProduct::try_new_array(lhs.clone(), rhs.clone())?.into_array(); + let original = + InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; let plugin = ScalarFnArrayPlugin::new(InnerProduct); let metadata = plugin diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index 7195265790f..7f04f8e7ed3 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -24,6 +24,7 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::builtins::ArrayBuiltins; @@ -428,7 +429,7 @@ pub fn normalize_as_l2_denorm( } // Calculate the norms of the vectors. - let norms_sfn = L2Norm::try_new_array(input.clone())?; + let norms_sfn = L2Norm.try_new_array(row_count, EmptyOptions, [input.clone()])?; let norms_array: ArrayRef = norms_sfn.into_array().execute(ctx)?; let primitive_norms: PrimitiveArray = norms_array.clone().execute(ctx)?; let norms_validity = primitive_norms.validity()?; diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 73c1538abc9..f2a96b37e23 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -13,7 +13,6 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; -use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; @@ -34,7 +33,6 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::serde::ArrayChildren; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -64,23 +62,6 @@ use crate::utils::validate_tensor_float_input; #[derive(Clone)] pub struct L2Norm; -impl L2Norm { - /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 norm operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(L2Norm, EmptyOptions) - } - - /// Constructs a [`ScalarFnArray`] that lazily computes the L2 norm over `child`. - /// - /// # Errors - /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). - pub fn try_new_array(child: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(L2Norm::new().erased(), vec![child]) - } -} - impl ScalarFnVTable for L2Norm { type Options = EmptyOptions; @@ -264,13 +245,14 @@ mod tests { use vortex_array::arrays::ConstantArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDType; use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_error::VortexResult; @@ -284,8 +266,7 @@ mod tests { /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. fn eval_l2_norm(input: ArrayRef) -> VortexResult> { - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; + let result = L2Norm.try_new_array(input.len(), EmptyOptions, [input])?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; Ok(prim.as_slice::().to_vec()) @@ -339,8 +320,7 @@ mod tests { let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; + let result = L2Norm.try_new_array(arr.len(), EmptyOptions, [arr])?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; @@ -359,8 +339,7 @@ mod tests { fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { let input = literal_vector_array(&[3.0f64, 4.0], 4); - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let result = L2Norm.try_new_array(input.len(), EmptyOptions, [input])?; let mut ctx = SESSION.create_execution_ctx(); let output = result.execute_until::(&mut ctx)?; @@ -390,8 +369,7 @@ mod tests { let null_scalar = Scalar::null(DType::Extension(ext_dtype)); let input = ConstantArray::new(null_scalar, 3).into_array(); - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let result = L2Norm.try_new_array(input.len(), EmptyOptions, [input])?; let mut ctx = SESSION.create_execution_ctx(); let output = result.execute_until::(&mut ctx)?; @@ -411,7 +389,7 @@ mod tests { #[case::fixed_shape_tensor(l2_norm_tensor_child())] #[case::vector(l2_norm_vector_child())] fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { - let original = L2Norm::try_new_array(child.clone())?.into_array(); + let original = L2Norm.try_new_array(child.len(), EmptyOptions, [child.clone()])?; let plugin = ScalarFnArrayPlugin::new(L2Norm); let metadata = plugin diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index ad3b96d1bff..492bc837b89 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -35,11 +35,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -79,7 +81,7 @@ pub fn build_similarity_search_tree>( let num_rows = data.len(); let query_vec = Vector::constant_array(query, num_rows)?; - let cosine = CosineSimilarity::try_new_array(data, query_vec)?.into_array(); + let cosine = CosineSimilarity.try_new_array(num_rows, EmptyOptions, [data, query_vec])?; let threshold_scalar = Scalar::primitive(threshold, Nullability::NonNullable); let threshold_array = ConstantArray::new(threshold_scalar, num_rows).into_array();