표현식 플러그인

표현식 플러그인

표현식 플러그인(expression plugins)은 사용자 정의 함수(UDF)를 만드는 데 권장되는 방법이에요. Rust 함수를 컴파일해 표현식으로 Polars 라이브러리에 등록할 수 있게 해 주죠. Polars 엔진이 런타임에 함수를 동적 링크하고, 여러분의 표현식은 네이티브 표현식과 거의 같은 속도로 실행됩니다. 게다가 Python 개입 없이 동작하기 때문에 GIL 경합도 없어요.

출처: 공식문서

플러그인은 기본 표현식이 누리는 것과 같은 이점을 그대로 얻습니다.

  • 최적화 (Optimization)
  • 병렬성 (Parallelism)
  • Rust 네이티브 성능

이제 커스텀 표현식을 만드는 데 필요한 것들을 알아볼게요.

첫 번째 커스텀 표현식: Pig Latin

첫 번째 표현식으로 pig latin 변환기를 만들어 볼게요. Pig Latin은 모든 단어에서 첫 글자를 떼어 뒤에 붙이고 마지막에 "ay"를 추가하는 우스운 언어예요. 그래서 "pig"라는 단어는 "igpay"로 변환됩니다.

물론 이미 표현식으로도 가능해요. 예를 들어 col("name").str.slice(1) + col("name").str.slice(0, 1) + "ay"처럼요. 하지만 이런 목적의 전용 함수가 더 잘 동작할 뿐 아니라, 플러그인을 배울 수 있게 해 줍니다.

준비하기

다음 Cargo.toml 파일로 새 라이브러리를 시작합니다.

[package]
name = "expression_lib"
version = "0.1.0"
edition = "2021"

[lib]
name = "expression_lib"
crate-type = ["cdylib"]

[dependencies]
polars = { version = "*" }
pyo3 = { version = "*", features = ["extension-module", "abi3-py310"] }
pyo3-polars = { version = "*", features = ["derive"] }
serde = { version = "*", features = ["derive"] }

표현식 작성하기

이 라이브러리에서 &str을 pig-latin으로 변환하는 헬퍼 함수를 만들고, 표현식으로 노출할 함수를 생성합니다. 함수를 노출하려면 #[polars_expr(output_type=DataType)] 어트리뷰트를 추가해야 하고, 함수는 항상 첫 번째 인자로 inputs: &[Series]를 받아야 해요.

// src/expressions.rs
use polars::prelude::*;
use pyo3_polars::derive::polars_expr;
use std::fmt::Write;

fn pig_latin_str(value: &str, output: &mut String) {
    if let Some(first_char) = value.chars().next() {
        write!(output, "{}{}ay", &value[1..], first_char).unwrap()
    }
}

#[polars_expr(output_type=String)]
fn pig_latinnify(inputs: &[Series]) -> PolarsResult<Series> {
    let ca = inputs[0].str()?;
    let out: StringChunked = ca.apply_into_string_amortized(pig_latin_str);
    Ok(out.into_series())
}

여기서 apply_values가 아니라 apply_into_string_amortized를 쓰는 이유는 행마다 새 문자열을 할당하지 않으려는 거예요. 플러그인이 여러 입력을 받고, 요소별(elementwise)로 동작하며, String 출력을 만든다면 polars::prelude::arity의 binary_elementwise_into_string_amortized 유틸리티 함수를 살펴보면 좋습니다.

이것이 Rust 쪽에서 필요한 전부예요. Python 쪽에서는 Cargo.toml에 정의한 이름과 같은 폴더를 만들어야 해요. 여기서는 "expression_lib"죠. Rust src 폴더와 같은 디렉터리에 expression_lib 폴더를 만들고 expression_lib/__init__.py를 생성합니다. 결과적인 파일 구조는 대략 이렇게 됩니다:

├── 📁 expression_lib/  # name must match "lib.name" in Cargo.toml
|   └── __init__.py
|
├── 📁src/
|   ├── lib.rs
|   └── expressions.rs
|
├── Cargo.toml
└── pyproject.toml

그다음 새 표현식들을 만듭니다. 표현식의 함수 이름을 등록할 수 있어요. 이 이름이 정확해야 한다는 점에 주의하세요. 그렇지 않으면 메인 Polars 패키지가 함수 이름을 해석할 수 없습니다. 또 Polars에 이 표현식이 어떻게 동작하는지 설명하는 추가 키워드 인자를 설정할 수 있어요. 이 경우 Polars에 이 함수가 요소별(elementwise)이라고 알려줍니다. 이렇게 하면 Polars가 이 표현식을 배치(batch)로 실행할 수 있어요. 반면 sort나 slice 같은 다른 연산에서는 이런 방식이 허용되지 않습니다.

# expression_lib/__init__.py
from pathlib import Path
from typing import TYPE_CHECKING

import polars as pl
from polars.plugins import register_plugin_function
from polars._typing import IntoExpr

PLUGIN_PATH = Path(__file__).parent

def pig_latinnify(expr: IntoExpr) -> pl.Expr:
    """Pig-latinnify expression."""
    return register_plugin_function(
        plugin_path=PLUGIN_PATH,
        function_name="pig_latinnify",
        args=expr,
        is_elementwise=True,
    )

그다음 maturin을 설치하고 maturin develop --release를 실행해 이 라이브러리를 환경에서 컴파일할 수 있어요.

그게 전부입니다. 우리의 표현식을 사용할 준비가 됐어요!

import polars as pl
from expression_lib import pig_latinnify

df = pl.DataFrame(
    {
        "convert": ["pig", "latin", "is", "silly"],
    }
)
out = df.with_columns(pig_latin=pig_latinnify("convert"))

대안으로 커스텀 네임스페이스를 등록할 수도 있는데, 그러면 Expr.language 네임스페이스를 만들어 사용자가 다음과 같이 쓸 수 있게 됩니다:

out = df.with_columns(
    pig_latin=pl.col("convert").language.pig_latinnify(),
)

kwargs 받기

Polars 표현식에서 kwargs(키워드 인자)를 받고 싶다면, Rust struct를 정의하고 serde::Deserialize를 파생시키기만 하면 돼요.

/// Provide your own kwargs struct with the proper schema and accept that type
/// in your plugin expression.
#[derive(Deserialize)]
pub struct MyKwargs {
    float_arg: f64,
    integer_arg: i64,
    string_arg: String,
    boolean_arg: bool,
}

/// If you want to accept `kwargs`. You define a `kwargs` argument
/// on the second position in you plugin. You can provide any custom struct that is deserializable
/// with the pickle protocol (on the Rust side).
#[polars_expr(output_type=String)]
fn append_kwargs(input: &[Series], kwargs: MyKwargs) -> PolarsResult<Series> {
    let input = &input[0];
    let input = input.cast(&DataType::String)?;
    let ca = input.str().unwrap();

    Ok(ca
        .apply_into_string_amortized(|val, buf| {
            write!(
                buf,
                "{}-{}-{}-{}-{}",
                val, kwargs.float_arg, kwargs.integer_arg, kwargs.string_arg, kwargs.boolean_arg
            )
                .unwrap()
        })
        .into_series())
}

Python 쪽에서는 플러그인을 등록할 때 kwargs를 넘길 수 있어요.

def append_args(
    expr: IntoExpr,
    float_arg: float,
    integer_arg: int,
    string_arg: str,
    boolean_arg: bool,
) -> pl.Expr:
    """
    This example shows how arguments other than `Series` can be used.
    """
    return register_plugin_function(
        plugin_path=PLUGIN_PATH,
        function_name="append_kwargs",
        args=expr,
        kwargs={
            "float_arg": float_arg,
            "integer_arg": integer_arg,
            "string_arg": string_arg,
            "boolean_arg": boolean_arg,
        },
        is_elementwise=True,
    )

출력 데이터 타입

출력 데이터 타입이 고정되어야 할 필요는 당연히 없어요. 종종 표현식의 입력 타입에 따라 달라지죠. 이를 처리하려면 #[polars_expr()] 매크로에 어떤 함수를 가리키는 output_type_func 인자를 제공하면 됩니다. 이 함수는 입력 필드 &[Field]를 출력 Field(이름과 데이터 타입)로 매핑할 수 있어요.

아래 스니펫은 이 매핑을 돕는 유틸리티 FieldsMapper를 사용하는 예시입니다.

use polars_plan::dsl::FieldsMapper;

fn haversine_output(input_fields: &[Field]) -> PolarsResult<Field> {
    FieldsMapper::new(input_fields).map_to_float_dtype()
}

#[polars_expr(output_type_func=haversine_output)]
fn haversine(inputs: &[Series]) -> PolarsResult<Series> {
    let out = match inputs[0].dtype() {
        DataType::Float32 => {
            let start_lat = inputs[0].f32().unwrap();
            let start_long = inputs[1].f32().unwrap();
            let end_lat = inputs[2].f32().unwrap();
            let end_long = inputs[3].f32().unwrap();
            crate::distances::naive_haversine(start_lat, start_long, end_lat, end_long)?
                .into_series()
        }
        DataType::Float64 => {
            let start_lat = inputs[0].f64().unwrap();
            let start_long = inputs[1].f64().unwrap();
            let end_lat = inputs[2].f64().unwrap();
            let end_long = inputs[3].f64().unwrap();
            crate::distances::naive_haversine(start_lat, start_long, end_lat, end_long)?
                .into_series()
        }
        _ => polars_bail!(InvalidOperation: "only supported for float types"),
    };
    Ok(out)
}

시작하는 데 알아야 할 것은 이게 전부예요. 이 모든 것이 어떻게 맞물리는지 보려면 이 저장소를, 더 깊이 있는 이해를 원하면 이 튜토리얼을 살펴보세요.

더 알아보기 (Learn more)