diff --git a/example/derive_expression/expression_lib/expression_lib/dist.py b/example/derive_expression/expression_lib/expression_lib/dist.py index 86e1e18..933d4a2 100644 --- a/example/derive_expression/expression_lib/expression_lib/dist.py +++ b/example/derive_expression/expression_lib/expression_lib/dist.py @@ -28,16 +28,15 @@ def jaccard_similarity(expr: IntoExpr, other: IntoExpr) -> pl.Expr: def haversine( - expr: IntoExpr, start_lat: IntoExpr, start_long: IntoExpr, end_lat: IntoExpr, end_long: IntoExpr, ) -> pl.Expr: - expr = parse_into_expr(expr) - return expr.register_plugin( + start_lat = parse_into_expr(start_lat) + return start_lat.register_plugin( lib=lib, - args=[start_lat, start_long, end_lat, end_long], + args=[start_long, end_lat, end_long], symbol="haversine", is_elementwise=True, cast_to_supertypes=True, diff --git a/example/derive_expression/run.py b/example/derive_expression/run.py index 49d8502..22094eb 100644 --- a/example/derive_expression/run.py +++ b/example/derive_expression/run.py @@ -10,7 +10,10 @@ "datetime": [datetime.now(tz=timezone.utc)] * 3, "dist_a": [[12, 32, 1], [], [1, -2]], "dist_b": [[-12, 1], [43], [876, -45, 9]], - "floats": [5.6, -1245.8, 242.224], + "start_lat": [5.6, -1245.8, 242.224], + "start_lon": [3.1, -1.1, 128.9], + "end_lat": [6.6, -1243.8, 240.224], + "end_lon": [3.9, -2, 130], } ) @@ -20,7 +23,7 @@ ).with_columns( hamming_dist=dist.hamming_distance("names", "pig_latin"), jaccard_sim=dist.jaccard_similarity("dist_a", "dist_b"), - haversine=dist.haversine("floats", "floats", "floats", "floats", "floats"), + haversine=dist.haversine("start_lat", "start_lon", "end_lat", "end_lon"), leap_year=date_util.is_leap_year("dates"), new_tz=date_util.change_time_zone("datetime"), appended_args=language.append_args( @@ -44,7 +47,7 @@ ).with_columns( hamming_dist=pl.col("names").dist.hamming_distance("pig_latin"), jaccard_sim=pl.col("dist_a").dist.jaccard_similarity("dist_b"), - haversine=pl.col("floats").dist.haversine("floats", "floats", "floats", "floats"), + haversine=pl.col("start_lat").dist.haversine("start_lon", "end_lat", "end_lon"), leap_year=pl.col("dates").date_util.is_leap_year(), new_tz=pl.col("datetime").date_util.change_time_zone(), appended_args=pl.col("names").language.append_args(