def test_array_return_type_reduction_window(con, t_random, df_random, qs):
expr = quantiles(t_random.b, qs).over(ibis.window())
result = expr.execute()
expected_raw = df_random.b.quantile(qs).tolist()
expected = pd.Series([expected_raw] * len(df_random))
expected = Spark.default_series_rename(expected)
tm.assert_series_equal(result, expected)
def test_udaf(con, t, df):
expr = my_string_length_sum(t.a)
assert isinstance(expr, ir.ScalarExpr)
result = expr.execute()
expected = t.a.execute().str.len().mul(2).sum()
expected = Spark.default_series_rename(expected)
assert result == expected
def test_udf(t, df, fn):
expr = fn(t.a)
assert isinstance(expr, ir.ColumnExpr)
result = expr.execute()
expected = df.a.str.len().mul(2)
expected = Spark.default_series_rename(expected)
tm.assert_series_equal(result, expected)
def test_multiple_argument_udf(con, t, df, fn):
expr = fn(t.b, t.c)
assert isinstance(expr, ir.ColumnExpr)
assert isinstance(expr, ir.NumericColumn)
assert isinstance(expr, ir.FloatingColumn)
result = expr.execute()
expected = df.b + df.c
expected = Spark.default_series_rename(expected)
tm.assert_series_equal(result, expected)
def test_compose_udfs(t_random, df_random, times_two_fn, add_one_fn):
expr = times_two_fn(add_one_fn(t_random.a))
result = expr.execute()
expected = df_random.a.add(1.0).mul(2.0)
expected = Spark.default_series_rename(expected)
tm.assert_series_equal(expected, result)