Vectorized UDF

示例

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

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

  import aura_frame as aura
  from aura_frame.udf import RegisterType
  from aura_frame.multimodal import ai_lake
  
  import pyarrow.compute as pc
  
  target_database = "test"
  # 隐式注册UDF
  @aura.udf.pyarrow(database=target_database, register_type=RegisterType.STAGED)
  def calculate_product(prices: aura.PyarrowVector[float], quantities: aura.PyarrowVector[int]) -> aura.PyarrowVector[float]:    
      return aura.PyarrowVector[float](pc.multiply(prices, quantities))
  
  # 使用UDF
  conn = 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
  )
  ds = conn.load_dataset("your-table", database=target_database)
  udf = conn.get_function("calculate_product", database=target_database)
  ds = ds.map(fn=udf, on=[ds.price, ds.quantity], as_col="product_column")
  ds.select_columns([ds.price, ds.quantity, ds.product_column])
  print(ds.execute())

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

  import aura_frame as aura
  from aura_frame.udf import RegisterType
  from aura_frame.multimodal import ai_lake
  
  import pandas as pd
  
  target_database = "test"
  # 隐式注册UDF
  @aura.udf.pandas(database=target_database, register_type=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())
  
  # 使用UDF
  conn = 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
  )
  ds = conn.load_dataset("your-table", database=target_database)
  udf = conn.get_function("calculate_product", database=target_database)
  ds = ds.map(fn=udf, on=[ds.price, ds.quantity], as_col="product_column")
  ds.select_columns([ds.price, ds.quantity, ds.product_column])
  print(ds.execute())