def test_sqrt_integer():
assert cudf.sqrt(16) == 4
assert_eq(cudf.sqrt(cudf.Series([4, 9, 16])), cudf.Series([2, 3, 4]))
assert_eq(
cudf.sqrt(cudf.DataFrame({"x": [4, 9, 16]})),
cudf.DataFrame({"x": [2, 3, 4]}),
)
def test_sqrt_float():
assert cudf.sqrt(16.0) == 4.0
assert_eq(cudf.sqrt(cudf.Series([4.0, 9, 16])), cudf.Series([2.0, 3, 4]))
assert_eq(
cudf.sqrt(cudf.DataFrame({"x": [4.0, 9, 16]})),
cudf.DataFrame({"x": [2.0, 3, 4]}),
)
def test_column_group_select():
df = cudf.DataFrame({
"a": [1, 4, 9, 16, 25],
"b": [0, 1, 2, 3, 4],
"c": [25, 16, 9, 4, 1]
})
input_features = ColumnGroup(["a", "b", "c"])
sqrt_features = input_features[["a", "c"]] >> cudf.sqrt
plus_one_features = input_features["b"] >> (lambda col: col + 1)
features = sqrt_features + plus_one_features
workflow = Workflow(features)
df_out = workflow.fit_transform(
Dataset(df)).to_ddf().compute(scheduler="synchronous")
expected = cudf.DataFrame()
expected["a"] = cudf.sqrt(df["a"])
expected["c"] = cudf.sqrt(df["c"])
expected["b"] = df["b"] + 1
assert_eq(expected, df_out)
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
if method == '__call__' and 'sqrt' == ufunc.__name__:
from cudf import sqrt
return sqrt(self)
else:
return NotImplemented