1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.text.clean_ip import CleanIp
10
11if __name__ == "__main__":
12    if os.getenv("DAFT_RUNNER", "native") == "ray":
13        import ray
14        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15        daft.set_runner_ray()
16    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18samples = {
19        "texts": [
20            "server at 192.168.1.1 is down",
21            "ipv6 2001:0db8:85a3:0000:0000:8a2e:0370:7334 unreachable",
22            "no ip here",
23            None,
24        ]
25    }
26    ds = daft.from_pydict(samples)
27    ds = ds.with_column(
28        "result",
29        aihc_udf(CleanIp, construct_args={}, num_cpus=1, concurrency=1, batch_size=4)(col("texts")),
30    )
31    ds.show()