Vectorized UDF

示例

下文提供两个示例展示如何使用Vectorized UDF。

• 示例一：使用PyArrow进行向量化加速。

  import os
  import aura_frame as aura
  from aura_frame.multimodal import ai_lake
  import pyarrow.compute as pc
  
  target_database = "test"
  
  def calculate_product(prices: aura.PyarrowVector[float], quantities: aura.PyarrowVector[int]) -> aura.PyarrowVector[float]:
      return aura.PyarrowVector[float](pc.multiply(prices, quantities))
  
  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,
  )
  
  try:
      udf = con.udf.pyarrow.register(
          calculate_product,
          name="calculate_product",
          database=target_database,
          register_type=aura.udf.RegisterType.STAGED,
      )
      ds = con.load_dataset("your-table", database=target_database)
      ds = ds.map(fn=udf, on=[ds.price, ds.quantity], as_col="product_column")
      ds = ds.select_columns(ds.price, ds.quantity, ds.product_column)
      print(ds.execute())
  finally:
      con.close()

• 示例二：使用Pandas进行向量化加速。

  import os
  import aura_frame as aura
  from aura_frame.multimodal import ai_lake
  import pandas as pd
  
  target_database = "test"
  # 隐式注册UDF
  @aura.udf.pandas(database=target_database, register_type=aura.udf.RegisterType.STAGED)
  def calculate_product(prices: aura.PandasVector[float], quantities: aura.PandasVector[int]) -> aura.PandasVector[float]:
      return aura.PandasVector[float](prices * quantities, dtype=pd.Float64Dtype())
  
  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,
  )
  
  try:
      ds = con.load_dataset("your-table", database=target_database)
      ds = ds.map(fn=calculate_product, on=[ds.price, ds.quantity], as_col="product_column")
      ds = ds.select_columns(ds.price, ds.quantity, ds.product_column)
      print(ds.execute())
  finally:
      con.close()