Go functions¶

When no expression says what you need, a Go function can: it runs inside the query, on every core, beside the expressions. A function call in Go is cheap, so a user-defined function costs far less here than a Python function does in Polars.

A value at a time: MapElements¶

MapElements(name, type, fn) calls fn on each value. The function's Go types are the column's in and the answer's out, inferred from fn by a generic method:

slug := ursus.Col("product").MapElements("slug", ursus.String,
    func(s string) (string, error) {
        return strings.ToUpper(s[:1]) + "-" + fmt.Sprint(len(s)), nil
    })
show(sales().Head(3).Select(ursus.Col("product"), slug.Alias("slug")))

Output

shape: (3, 2)
┌─────────┬────────┐
│ product │ slug   │
│ String  │ String │
├─────────┼────────┤
│ "apple" │ "A-5"  │
│ "apple" │ "A-5"  │
│ "pear"  │ "P-4"  │
└─────────┴────────┘

An answer can be of any type, and a mapped column is an expression like any other, a group-by's key here:

band := ursus.Col("qty").MapElements("band", ursus.String,
    func(q int64) (string, error) {
        switch {
        case q >= 10:
            return "large", nil
        case q >= 5:
            return "medium", nil
        }
        return "small", nil
    })
show(sales().
    GroupBy(band.Alias("band")).
    Agg(ursus.Len().Alias("orders")).
    Sort(ursus.Asc(ursus.Col("band"))))

Output

shape: (3, 2)
┌──────────┬────────┐
│ band     │ orders │
│ String   │ Uint64 │
├──────────┼────────┤
│ "large"  │ 2      │
│ "medium" │ 3      │
│ "small"  │ 3      │
└──────────┴────────┘

Nulls pass through: fn is never called with one, so it cannot mistake a missing value for a zero. Three of the six rows were null, and the function ran three times:

var calls atomic.Int64
double := ursus.Col("temp").MapElements("double", ursus.Float64,
    func(t float64) (float64, error) {
        calls.Add(1)
        return 2 * t, nil
    })
show(readings().Select(ursus.Col("temp"), double.Alias("doubled")))
fmt.Println("calls:", calls.Load())

Output

shape: (6, 2)
┌─────────┬─────────┐
│ temp    │ doubled │
│ Float64 │ Float64 │
├─────────┼─────────┤
│ 10      │ 20      │
│ null    │ null    │
│ null    │ null    │
│ 13      │ 26      │
│ 12      │ 24      │
│ null    │ null    │
└─────────┴─────────┘

calls: 3

A column at a time: MapBatches¶

MapBatches(name, type, fn) calls fn with a batch of values and whether each is there, and takes back the same: the place to decide what a null means, or to call something that works on many values at once.

orZero := ursus.Col("temp").MapBatches("or_zero", ursus.Float64,
    func(vals []float64, valid []bool) ([]float64, []bool, error) {
        out, ok := make([]float64, len(vals)), make([]bool, len(vals))
        for i := range vals {
            if valid[i] {
                out[i] = vals[i]
            }
            ok[i] = true // a null becomes 0, a value of its own
        }
        return out, ok, nil
    })
show(readings().Select(ursus.Col("temp"), orZero.Alias("or_zero")))

Output

shape: (6, 2)
┌─────────┬─────────┐
│ temp    │ or_zero │
│ Float64 │ Float64 │
├─────────┼─────────┤
│ 10      │ 10      │
│ null    │ 0       │
│ null    │ 0       │
│ 13      │ 13      │
│ 12      │ 12      │
│ null    │ 0       │
└─────────┴─────────┘

A window at a time: RollingMap¶

New in 0.6.7. RollingMap(name, type, size, fn) calls fn on each row's window of size rows, as Polars' rolling_map: the window's values in its order and whether each is there, as MapBatches gives a batch, and the answer is the row's. A window of fewer values than MinSamples(k), size unless set, is null; RollingCenter() centres it on its row. It runs within each partition under Over, and streams past a memory limit as the rolling functions do.

A row at a time: MapRows¶

New in 0.6.7. A frame's MapRows(ctx, fn) calls fn on each row, as Polars' map_rows. A row is a struct, read as Rows reads one, or a []any of its values in the columns' order; a struct answer is a column a field, and any other answer one column named map. The rows are decoded and fn called once each, in order, on one goroutine, so an expression or MapBatches is the faster way where one can say it.

A whole column: MapColumns¶

New in 0.6.7. A frame's MapColumns(ctx, names, fn) replaces each named column with fn of it, as Polars' map_columns: the answer keeps the column's name and place, and is as long as the frame, or one value, which every row takes.

Errors¶

An error fn returns fails the query, naming the function, the value and the row of its batch. A panic in fn is reported the same way, rather than ending the program.

dozens := ursus.Col("qty").MapElements("dozens", ursus.Float64,
    func(q int64) (float64, error) {
        if q > 10 {
            return 0, errors.New("more than ten is a wholesale order")
        }
        return float64(q) / 12, nil
    })
_, err := sales().Select(dozens).Collect(context.Background())
fmt.Println(err)

Output

ursus: map_elements: udf "dozens" failed on 12, row 5 of its batch of column "qty"
  caused by: more than ten is a wholesale order

The name¶

The name is how Explain and every error show a function, since a closure prints as nothing. Two different functions given one name are refused before the query runs:

slug := ursus.Col("product").MapElements("slug", ursus.String,
    func(s string) (string, error) { return strings.ToUpper(s), nil })
plan, err := sales().
    Filter(ursus.Col("qty").Gt(5)).
    Select(slug).
    Explain(context.Background())
if err != nil {
    log.Fatal(err)
}
fmt.Println(plan)

Output

PROJECT [col("product").map_elements("slug" -> String)]
  FILTER [(col("qty") > lit(5))]
    CSV SCAN testdata/sales.csv
      projection: [product, qty] (2/6 cols)

A function used twice should be built once and held in a variable, so it is one function under one name.

What the function must do¶

Be safe to call from several goroutines. A query runs on every core, and there is no opting one expression out: a closure over a shared map is a data race. Count with sync/atomic, as the null example does.
Not depend on the order it is called in. The answer's order is the rows', the calls' is not.
Answer the type it declares. The declared type is the column's from the plan's start; an answer of another type fails the query.