Skip to content

Use Python packages and plugins

Polars Nanook runs expression plugins and Python user-defined functions (UDFs) in remote queries. The packages they need must be installed on the workers, and in the client's environment too, because the client imports them to build the query.

Set up

List the packages in a requirements.txt file:

requirements.txt
numpy==2.2.4
polars-xdt==0.16.8

Then pass the file in the requirements argument of the compute context:

ComputeContext

import polars_nanook as pn

ctx = pn.ComputeContext(
    workspace="My workspace",
    cpus=12,
    memory=12,
    requirements="requirements.txt",
)

When this context starts, Polars Nanook installs the packages on the compute cluster before running queries.

List the packages in runtime.composed.requirements, or build them into an image set with runtime.prebuilt. See Worker image.

Install the packages in the Python environment that runs polars-nanook-compute on every worker machine. See Python environment.

Expression plugins in a remote query

from datetime import datetime

import polars as pl
import polars_nanook as pn
import polars_xdt as xdt

ctx: pn.ComputeContext | pn.ClusterContext

lf = pl.LazyFrame(
    {
        "local_dt": [
            datetime(2020, 10, 10, 1),
            datetime(2020, 10, 10, 2),
            datetime(2020, 10, 9, 20),
        ],
        "timezone": [
            "Europe/London",
            "Africa/Kigali",
            "America/New_York",
        ],
    }
)

query = lf.with_columns(
    xdt.from_local_datetime("local_dt", pl.col("timezone"), "UTC").alias("date")
)

query.remote(ctx).collect()

Python UDFs in a remote query

The client pickles each UDF, so on a cluster you run with the Helm chart or on bare-metal, the client needs the cluster's Python release (Set up).

import numpy as np
import polars as pl
import polars_nanook as pn

ctx: pn.ComputeContext | pn.ClusterContext

lf = pl.LazyFrame(
    {
        "keys": ["a", "a", "b", "b"],
        "values": [10, 7, 1, 23],
    }
)

q = lf.select(
    pl.col("values").map_batches(np.log, return_dtype=pl.Float64, is_elementwise=True)
)

q.remote(ctx).collect()