Skip to content

Commit 4363c00

Browse files
committed
Standardize scalar function array construction
Uses `ScalarFnFactoryExt` for tensor and geo scalar functions instead of maintaining per-function constructors. `L2Denorm` keeps its checked constructor because it validates the normalized-row invariant. Signed-off-by: "Connor Tsui" <connor.tsui20@gmail.com>
1 parent 2de0319 commit 4363c00

11 files changed

Lines changed: 97 additions & 187 deletions

File tree

‎vortex-geo/benches/envelope.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ use vortex_array::Canonical;
2525
use vortex_array::ExecutionCtx;
2626
use vortex_array::IntoArray;
2727
use vortex_array::VortexSessionExecute;
28+
use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
29+
use vortex_array::scalar_fn::EmptyOptions;
2830
use vortex_geo::scalar_fn::envelope::GeoEnvelope;
2931
use vortex_geo::test_harness::MultiPolygonRings;
3032
use vortex_geo::test_harness::geo_session;
@@ -58,9 +60,9 @@ fn coin(i: usize) -> bool {
5860

5961
/// Execute the envelope of `column` to completion.
6062
fn envelope(column: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef {
61-
GeoEnvelope::try_new_array(column.clone())
63+
GeoEnvelope
64+
.try_new_array(column.len(), EmptyOptions, [column.clone()])
6265
.unwrap()
63-
.into_array()
6466
.execute::<Canonical>(ctx)
6567
.unwrap()
6668
.into_array()

‎vortex-geo/src/scalar_fn/contains.rs‎

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
use geo::Contains;
77
use vortex_array::ArrayRef;
88
use vortex_array::ExecutionCtx;
9-
use vortex_array::arrays::ScalarFnArray;
109
use vortex_array::dtype::DType;
1110
use vortex_array::dtype::Nullability;
1211
use vortex_array::expr::Expression;
@@ -17,7 +16,6 @@ use vortex_array::scalar_fn::EmptyOptions;
1716
use vortex_array::scalar_fn::ExecutionArgs;
1817
use vortex_array::scalar_fn::ScalarFnId;
1918
use vortex_array::scalar_fn::ScalarFnVTable;
20-
use vortex_array::scalar_fn::TypedScalarFnInstance;
2119
use vortex_error::VortexResult;
2220
use vortex_session::VortexSession;
2321
use vortex_session::registry::CachedId;
@@ -31,17 +29,6 @@ use crate::scalar_fn::execute::execute_null_propagating;
3129
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
3230
pub struct GeoContains;
3331

34-
impl GeoContains {
35-
/// A lazy `ScalarFnArray` computing per-row whether operand `a` contains operand `b`;
36-
/// either may be constant. The output length is taken from `a`.
37-
pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult<ScalarFnArray> {
38-
ScalarFnArray::try_new(
39-
TypedScalarFnInstance::new(GeoContains, EmptyOptions).erased(),
40-
vec![a, b],
41-
)
42-
}
43-
}
44-
4532
impl ScalarFnVTable for GeoContains {
4633
type Options = EmptyOptions;
4734

@@ -127,6 +114,7 @@ mod tests {
127114
use vortex_array::VortexSessionExecute;
128115
use vortex_array::arrays::BoolArray;
129116
use vortex_array::arrays::ConstantArray;
117+
use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
130118
use vortex_array::assert_arrays_eq;
131119
use vortex_array::dtype::DType;
132120
use vortex_array::dtype::Nullability;
@@ -177,7 +165,7 @@ mod tests {
177165
) -> VortexResult<()> {
178166
let session = vortex_array::array_session();
179167
let mut ctx = session.create_execution_ctx();
180-
let contains = GeoContains::try_new_array(a, b)?.into_array();
168+
let contains = GeoContains.try_new_array(a.len(), EmptyOptions, [a, b])?;
181169
assert_arrays_eq!(contains, BoolArray::from_iter(expected), &mut ctx);
182170
Ok(())
183171
}
@@ -298,7 +286,8 @@ mod tests {
298286

299287
let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?;
300288
let points = nullable_point_column(vec![Some((2.0, 2.0)), None, Some((10.0, 10.0))])?;
301-
let contains = GeoContains::try_new_array(container, points)?.into_array();
289+
let contains =
290+
GeoContains.try_new_array(container.len(), EmptyOptions, [container, points])?;
302291

303292
let expected = BoolArray::new(
304293
BitBuffer::from_iter([true, false, false]),
@@ -318,7 +307,8 @@ mod tests {
318307
let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable();
319308
let null_const = ConstantArray::new(Scalar::null(point_dtype), 2).into_array();
320309
let points = point_column(vec![2.0, 10.0], vec![2.0, 10.0])?;
321-
let contains = GeoContains::try_new_array(null_const, points)?.into_array();
310+
let contains =
311+
GeoContains.try_new_array(null_const.len(), EmptyOptions, [null_const, points])?;
322312

323313
let expected =
324314
BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array();
@@ -346,7 +336,8 @@ mod tests {
346336
None,
347337
Some((4.0, 4.0)),
348338
])?;
349-
let contains = GeoContains::try_new_array(container, contained)?.into_array();
339+
let contains =
340+
GeoContains.try_new_array(container.len(), EmptyOptions, [container, contained])?;
350341

351342
let expected = BoolArray::new(
352343
BitBuffer::from_iter([true, false, false, false]),
@@ -365,7 +356,8 @@ mod tests {
365356

366357
let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?;
367358
let points = nullable_point_column(vec![None, None])?;
368-
let contains = GeoContains::try_new_array(container, points)?.into_array();
359+
let contains =
360+
GeoContains.try_new_array(container.len(), EmptyOptions, [container, points])?;
369361

370362
let expected =
371363
BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array();
@@ -382,7 +374,8 @@ mod tests {
382374

383375
let container = nullable_point_column(vec![Some((1.0, 1.0)), None])?;
384376
let contained = nullable_point_column(vec![None, Some((2.0, 2.0))])?;
385-
let contains = GeoContains::try_new_array(container, contained)?.into_array();
377+
let contains =
378+
GeoContains.try_new_array(container.len(), EmptyOptions, [container, contained])?;
386379

387380
let expected =
388381
BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array();

‎vortex-geo/src/scalar_fn/distance.rs‎

Lines changed: 14 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ use geo::Distance;
77
use geo::Euclidean;
88
use vortex_array::ArrayRef;
99
use vortex_array::ExecutionCtx;
10-
use vortex_array::arrays::ScalarFnArray;
1110
use vortex_array::dtype::DType;
1211
use vortex_array::dtype::Nullability;
1312
use vortex_array::dtype::PType;
@@ -19,7 +18,6 @@ use vortex_array::scalar_fn::EmptyOptions;
1918
use vortex_array::scalar_fn::ExecutionArgs;
2019
use vortex_array::scalar_fn::ScalarFnId;
2120
use vortex_array::scalar_fn::ScalarFnVTable;
22-
use vortex_array::scalar_fn::TypedScalarFnInstance;
2321
use vortex_error::VortexResult;
2422
use vortex_session::VortexSession;
2523
use vortex_session::registry::CachedId;
@@ -32,17 +30,6 @@ use crate::scalar_fn::execute::execute_null_propagating;
3230
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
3331
pub struct GeoDistance;
3432

35-
impl GeoDistance {
36-
/// A lazy `ScalarFnArray` computing the per-row distance between operands `a` and `b`; either may
37-
/// be constant. The output length is taken from `a`.
38-
pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult<ScalarFnArray> {
39-
ScalarFnArray::try_new(
40-
TypedScalarFnInstance::new(GeoDistance, EmptyOptions).erased(),
41-
vec![a, b],
42-
)
43-
}
44-
}
45-
4633
impl ScalarFnVTable for GeoDistance {
4734
type Options = EmptyOptions;
4835

@@ -115,6 +102,7 @@ mod tests {
115102
use vortex_array::VortexSessionExecute;
116103
use vortex_array::arrays::ConstantArray;
117104
use vortex_array::arrays::PrimitiveArray;
105+
use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
118106
use vortex_array::assert_arrays_eq;
119107
use vortex_array::dtype::DType;
120108
use vortex_array::dtype::Nullability;
@@ -159,7 +147,7 @@ mod tests {
159147

160148
let a = point_column(vec![0.0, 3.0, 0.0, 3.0], vec![0.0, 0.0, 4.0, 4.0])?;
161149
let b = point_constant(0.0, 0.0, 4, &mut ctx)?;
162-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
150+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
163151

164152
assert_eq!(distances(distance, &mut ctx)?, vec![0.0, 3.0, 4.0, 5.0]);
165153
Ok(())
@@ -173,7 +161,7 @@ mod tests {
173161

174162
let a = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?;
175163
let b = point_column(vec![3.0, 1.0], vec![4.0, 1.0])?;
176-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
164+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
177165

178166
assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 0.0]);
179167
Ok(())
@@ -190,7 +178,7 @@ mod tests {
190178
let single = polygon_column(vec![vec![ring]])?.execute_scalar(0, &mut ctx)?;
191179
let square = ConstantArray::new(single, 2).into_array();
192180
let points = point_column(vec![7.0, 2.0], vec![2.0, 2.0])?;
193-
let distance = GeoDistance::try_new_array(points, square)?.into_array();
181+
let distance = GeoDistance.try_new_array(points.len(), EmptyOptions, [points, square])?;
194182

195183
assert_eq!(distances(distance, &mut ctx)?, vec![3.0, 0.0]);
196184
Ok(())
@@ -204,7 +192,7 @@ mod tests {
204192

205193
let a = point_constant(0.0, 0.0, 4, &mut ctx)?;
206194
let b = point_column(vec![0.0, 3.0, 0.0, 3.0], vec![0.0, 0.0, 4.0, 4.0])?;
207-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
195+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
208196

209197
assert_eq!(distances(distance, &mut ctx)?, vec![0.0, 3.0, 4.0, 5.0]);
210198
Ok(())
@@ -218,7 +206,7 @@ mod tests {
218206

219207
let a = point_constant(0.0, 0.0, 3, &mut ctx)?;
220208
let b = point_constant(3.0, 4.0, 3, &mut ctx)?;
221-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
209+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
222210

223211
assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 5.0, 5.0]);
224212
Ok(())
@@ -245,7 +233,7 @@ mod tests {
245233

246234
let a = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((3.0, 4.0))])?;
247235
let b = point_constant(0.0, 0.0, 3, &mut ctx)?;
248-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
236+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
249237

250238
let expected = PrimitiveArray::new(
251239
vec![0.0f64, 0.0, 5.0],
@@ -264,7 +252,7 @@ mod tests {
264252

265253
let a = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((0.0, 0.0))])?;
266254
let b = nullable_point_column(vec![Some((3.0, 4.0)), Some((1.0, 1.0)), None])?;
267-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
255+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
268256

269257
let expected = PrimitiveArray::new(
270258
vec![5.0f64, 0.0, 0.0],
@@ -284,7 +272,8 @@ mod tests {
284272
let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable();
285273
let null_const = ConstantArray::new(Scalar::null(point_dtype), 3).into_array();
286274
let b = point_column(vec![0.0, 3.0, 0.0], vec![0.0, 0.0, 4.0])?;
287-
let distance = GeoDistance::try_new_array(null_const, b)?.into_array();
275+
let distance =
276+
GeoDistance.try_new_array(null_const.len(), EmptyOptions, [null_const, b])?;
288277

289278
let expected = PrimitiveArray::new(vec![0.0f64; 3], Validity::AllInvalid).into_array();
290279
assert_arrays_eq!(distance, expected, &mut ctx);
@@ -299,7 +288,7 @@ mod tests {
299288

300289
let a = nullable_point_column(vec![None, None])?;
301290
let b = point_constant(0.0, 0.0, 2, &mut ctx)?;
302-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
291+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
303292

304293
let expected = PrimitiveArray::new(vec![0.0f64; 2], Validity::AllInvalid).into_array();
305294
assert_arrays_eq!(distance, expected, &mut ctx);
@@ -315,7 +304,7 @@ mod tests {
315304

316305
let a = nullable_point_column(vec![Some((0.0, 0.0)), None])?;
317306
let b = nullable_point_column(vec![None, Some((1.0, 1.0))])?;
318-
let distance = GeoDistance::try_new_array(a, b)?.into_array();
307+
let distance = GeoDistance.try_new_array(a.len(), EmptyOptions, [a, b])?;
319308

320309
let expected = PrimitiveArray::new(vec![0.0f64; 2], Validity::AllInvalid).into_array();
321310
assert_arrays_eq!(distance, expected, &mut ctx);
@@ -332,8 +321,8 @@ mod tests {
332321

333322
let a = point_column(vec![], vec![])?;
334323
let b = point_column(vec![], vec![])?;
335-
let result = GeoDistance::try_new_array(a, b)?
336-
.into_array()
324+
let result = GeoDistance
325+
.try_new_array(a.len(), EmptyOptions, [a, b])?
337326
.execute::<Canonical>(&mut ctx)?
338327
.into_array();
339328

‎vortex-geo/src/scalar_fn/envelope.rs‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ use vortex_array::ArrayRef;
1010
use vortex_array::ExecutionCtx;
1111
use vortex_array::IntoArray;
1212
use vortex_array::arrays::ExtensionArray;
13-
use vortex_array::arrays::ScalarFnArray;
1413
use vortex_array::arrays::StructArray;
1514
use vortex_array::arrays::extension::ExtensionArrayExt;
1615
use vortex_array::arrays::struct_::StructArrayExt;
@@ -25,7 +24,6 @@ use vortex_array::scalar_fn::EmptyOptions;
2524
use vortex_array::scalar_fn::ExecutionArgs;
2625
use vortex_array::scalar_fn::ScalarFnId;
2726
use vortex_array::scalar_fn::ScalarFnVTable;
28-
use vortex_array::scalar_fn::TypedScalarFnInstance;
2927
use vortex_array::validity::Validity;
3028
use vortex_buffer::BitBuffer;
3129
use vortex_buffer::BufferMut;
@@ -53,17 +51,6 @@ use crate::extension::validate_geometry_operands;
5351
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
5452
pub struct GeoEnvelope;
5553

56-
impl GeoEnvelope {
57-
/// A lazy `ScalarFnArray` computing the per-row bounding box of geometry operand `a`, which may
58-
/// be constant. The output length is taken from `a`.
59-
pub fn try_new_array(a: ArrayRef) -> VortexResult<ScalarFnArray> {
60-
ScalarFnArray::try_new(
61-
TypedScalarFnInstance::new(GeoEnvelope, EmptyOptions).erased(),
62-
vec![a],
63-
)
64-
}
65-
}
66-
6754
/// The output dtype: a nullable native 2-D box ([`Rect`], `geoarrow.box`) column. Nullable
6855
/// because rows without a box — null or empty geometries — are null. Metadata is defaulted.
6956
fn output_box_dtype() -> VortexResult<ExtDType<Rect>> {
@@ -231,6 +218,7 @@ mod tests {
231218
use vortex_array::arrays::ExtensionArray;
232219
use vortex_array::arrays::PrimitiveArray;
233220
use vortex_array::arrays::StructArray;
221+
use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
234222
use vortex_array::assert_arrays_eq;
235223
use vortex_array::dtype::DType;
236224
use vortex_array::dtype::Nullability;
@@ -259,7 +247,7 @@ mod tests {
259247

260248
/// Execute a `GeoEnvelope` over `array`, returning the lazy box column.
261249
fn boxes(array: ArrayRef) -> VortexResult<ArrayRef> {
262-
Ok(GeoEnvelope::try_new_array(array)?.into_array())
250+
GeoEnvelope.try_new_array(array.len(), EmptyOptions, [array])
263251
}
264252

265253
/// A point's box is degenerate: both corners are the point itself.

0 commit comments

Comments
 (0)