UDTF

示例

以下示例展示UDTF的注册和使用的流程。

1. 创建会话。

   import ibis
   import aura_frame as aura
   import logging
   from aura_frame.multimodal import ai_lake
   
   target_database = "test"
   # create session
   con = ai_lake.connect(
           aura_endpoint=os.getenv("aura_endpoint"),
           aura_endpoint_name=os.getenv("aura_endpoint_name"),
           aura_workspace_id=os.getenv("aura_workspace_id"),
           lf_catalog_name=os.getenv("lf_catalog_name"),
           access_key=os.getenv("access_key"),
           secret_key=os.getenv("secret_key"),
           default_database=target_database,
           use_single_cn_mode=True
   )

2. 定义并注册UDTF。

   # function definition
   def py_generate_series(start, endnum, step):
       current = start
       while current <= endnum:
           yield {'col1':current}
           current += step
   
   # register udtf
   import ibis.expr.datatypes as dt
   func_signature=aura.Signature(
        parameters=[
            aura.Parameter(name="start", annotation=int),
            aura.Parameter(name="endnum", annotation=int),
            aura.Parameter(name="step", annotation=int)
        ],
        return_annotation=dt.Struct({"col1": int}),
   )
   target_database = 'test'
   conncreate_table_function(py_generate_series, database=target_database, signature=func_signature)

3. 调用UDTF。

   # use udtf
   ds = connload_dataset("demo_table", database=target_database)
   py_generate_series_handler = connget_function("py_generate_series", database=target_database)
   ds = ds.flat_map(
      fn=py_generate_series_handler,
      on=[ds.a, ds.b, ds.c],
      as_col='udtf_col'
   )
   
   # trigger executing
   res = ds.execute()
   print(res)
   | a     | b     | c     | udtf_col |
   | ----- | ----- | ----- | -------- |
   | int64 | int64 | int64 | int64    |
   | ----- | ----- | ----- | -------- |
   | 1     | 4     | 1     | 1        |
   | 1     | 4     | 1     | 2        |
   | 1     | 4     | 1     | 3        |
   | 1     | 4     | 1     | 4        |