circle_func.py

@udf.jit()
def filter_func(out: udf.Array(udf.float32, 2),
                vals: udf.Array(udf.float32, 2),
                nrows: udf.int64, ncols: udf.int64,
                iscircle: udf.bool) -> udf.int32:
    n = out.window_shape[0]
    m = out.window_shape[1]
    row_start = out.window_start[0]
    col_start = out.window_start[1]
    for i in range(n):
        for j in range(m):
            if iscircle:
                out[i, j] = ulib.circle_filter(
                    vals[i, j], row_start + i, col_start + j,
                    nrows, ncols)
            else:
                out[i, j] = ulib.square_filter(vals[i, j])
    return 0