Vectorized UDF

示例

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

• 示例一：使用PyArrow进行向量化加速。通过 `con.create_scalar_function` 显式注册PyArrow向量化UDF，框架会根据类型注解自动识别为PyArrow类型，再通过`ds.map`调用。

  import os
  import ibis
  import aura_frame as aura
  from aura_frame.multimodal import ai_lake
  from aura_frame.multimodal.types.vectorized import PyarrowVector
  import pyarrow.compute as pc
  
  target_database = "test"  #在LakeFormation中创建数据库，例如：test。具体操作，请参见创建LakeFormation实例并规划元数据
  #连接AuraJobV2端点，具体参数说明，请参见建立与AuraJobV2端点的连接。
  con = ai_lake.connect(
      aura_endpoint=os.getenv("aura_endpoint"),         # AuraJobV2端点所在的终端节点（Endpoint）
      aura_endpoint_id=os.getenv("aura_endpoint_id"),    # AuraJobV2端点ID
      aura_endpoint_name=os.getenv("aura_endpoint_name"),# AuraJobV2端点名称
      aura_workspace_id=os.getenv("aura_workspace_id"),  # 工作空间ID
      lf_catalog_name=os.getenv("lf_catalog_name"),      # LakeFormation Catalog名称
      access_key=os.getenv("access_key"),                # 访问密钥ID（AK）
      secret_key=os.getenv("secret_key"),                # 访问密钥（SK）
      default_database="postgres",                       # LakeFormation数据库名称
      use_single_cn_mode=True,   # 是否开启单CN模式
      logging_level=logging.INFO # 设置日志级别
  )
  #配置UDF部署工作区，具体参数说明，请参见配置UDF部署工作区。
  con.set_function_staging_workspace(
      obs_directory_base=os.getenv("obs_directory_base"),  # 并行文件系统中UDF的存储路径
      obs_bucket_name=os.getenv("obs_bucket_name"),       # 并行文件系统名称
      obs_server=os.getenv("obs_server"),                  # OBS终端节点（Endpoint）
      access_key=os.getenv("access_key"),                  # 访问密钥ID（AK）
      secret_key=os.getenv("secret_key"),                  # 访问密钥（SK）
  )
  def calculate_product_pyarrow(
      prices: PyarrowVector[float],
      quantities: PyarrowVector[int],
  ) -> PyarrowVector[float]:
      return PyarrowVector[float](pc.multiply(prices, quantities))
  
  try:
      # 创建测试表并写入数据
      table_name = "sales_records"
      table_schema = ibis.schema({"price": "float64", "quantity": "int32"})
      con.sql(f'DROP TABLE IF EXISTS "{target_database}"."{table_name}";')
      con.create_table(
          table_name,
          schema=table_schema,
          database=target_database,
          table_format="orc",
          if_not_exists=True,
      )
      con.sql(f'''
      INSERT INTO "{target_database}"."{table_name}" (price, quantity) VALUES
          (100.0, 10),
          (200.0, 20),
          (150.0, 15)
      ''')
      # 删除已存在的同名UDF（支持重复执行）
      con.delete_function("calculate_product_pyarrow", database=target_database, if_it_exists=True)
      # 显式注册 PyArrow 向量化 UDF
      udf_pyarrow = con.create_scalar_function(
          calculate_product_pyarrow,
          name="calculate_product_pyarrow",
          database=target_database,
          register_type=aura.udf.RegisterType.STAGED,
      )
      ds = con.load_dataset("sales_records", database=target_database)
      ds = ds.map(fn=udf_pyarrow, on=[ds.price, ds.quantity], as_col="product_column")
      ds = ds.select_columns(ds.price, ds.quantity, ds.product_column)
      print("PyArrow 向量化结果:")
      print(ds.execute())
  finally:
      con.close()

• 示例二：使用Pandas进行向量化加速。通过`@aura.udf.pandas`装饰器隐式注册Pandas向量化UDF，注册类型为`STAGED`。

  import os
  import ibis
  import aura_frame as aura
  from aura_frame.multimodal import ai_lake
  from aura_frame.multimodal.types.vectorized import PandasVector
  import pandas as pd
  
  target_database = "test"
  # 隐式注册UDF
  @aura.udf.pandas(database=target_database, register_type=aura.udf.RegisterType.STAGED)
  def calculate_product_pandas(
      prices: PandasVector[float],
      quantities: PandasVector[int],
  ) -> PandasVector[float]:
      return PandasVector[float](prices * quantities, dtype=pd.Float64Dtype())
  con = ai_lake.connect(
      aura_endpoint=os.getenv("aura_endpoint"),         # AuraJobV2端点所在的终端节点（Endpoint）
      aura_endpoint_id=os.getenv("aura_endpoint_id"),    # AuraJobV2端点ID
      aura_endpoint_name=os.getenv("aura_endpoint_name"),# AuraJobV2端点名称
      aura_workspace_id=os.getenv("aura_workspace_id"),  # 工作空间ID
      lf_catalog_name=os.getenv("lf_catalog_name"),      # LakeFormation Catalog名称
      lf_instance_id=os.getenv("lf_instance_id"),        # LakeFormation实例ID
      project_id=os.getenv("project_id"),                # 项目ID
      access_key=os.getenv("access_key"),                # 访问密钥ID（AK）
      secret_key=os.getenv("secret_key"),                # 访问密钥（SK）
      default_database=target_database,                  # 默认数据库
      use_single_cn_mode=True,                           # 是否开启单CN模式
  )
  con.set_function_staging_workspace(
      obs_directory_base=os.getenv("obs_directory_base"),  # 并行文件系统中UDF的存储路径
      obs_bucket_name=os.getenv("obs_bucket_name"),        # 并行文件系统名称
      obs_server=os.getenv("obs_server"),                  # OBS终端节点（Endpoint）
      access_key=os.getenv("access_key"),                  # 访问密钥ID（AK）
      secret_key=os.getenv("secret_key"),                  # 访问密钥（SK）
  )
  try:
      # 创建测试表并写入数据
      table_name = "sales_records"
      table_schema = ibis.schema({"price": "float64", "quantity": "int32"})
      con.sql(f'DROP TABLE IF EXISTS "{target_database}"."{table_name}";')
      con.create_table(
          table_name,
          schema=table_schema,
          database=target_database,
          table_format="orc",
          if_not_exists=True,
      )
      con.sql(f'''
      INSERT INTO "{target_database}"."{table_name}" (price, quantity) VALUES
          (100.0, 10),
          (200.0, 20),
          (150.0, 15)
      ''')
      # 删除已存在的同名UDF（支持重复执行）
      con.delete_function("calculate_product_pandas", database=target_database, if_it_exists=True)
      # 使用 Pandas 向量化 UDF（已在上方通过装饰器隐式注册）
      ds2 = con.load_dataset("sales_records", database=target_database)
      ds2 = ds2.map(fn=calculate_product_pandas, on=[ds2.price, ds2.quantity], as_col="product_column")
      ds2 = ds2.select_columns(ds2.price, ds2.quantity, ds2.product_column)
      print("Pandas 向量化结果:")
      print(ds2.execute())
  finally:
      con.close()