Build a data pipeline

Filter, sort and join two datasets end to end

This tutorial builds a small analytics pipeline over two related datasets — employees and the departments they belong to. Along the way you’ll use array construction, compute kernels, row filtering, sorting and a hash join. Every cell runs when the docs are built.

The data

We start with two RecordBatches. The first holds employees; the second maps department ids to names.

employees = ma.record_batch({
    "name":    ma.array(["Alice", "Bob", "Carol", "Dave", "Eve", "Frank"]),
    "dept_id": ma.array([1, 2, 1, 3, 2, 1]),
    "salary":  ma.array([95_000, 72_000, 120_000, 88_000, 61_000, 105_000]),
    "age":     ma.array([34, 28, 45, 52, 24, 39]),
})

departments = ma.record_batch({
    "dept_id": ma.array([1, 2, 3]),
    "dept":    ma.array(["Engineering", "Sales", "Marketing"]),
})

print(employees)
print(departments)
RecordBatch(num_rows=6, schema=Schema(fields=[name: string, dept_id: int64, salary: int64, age: int64]))
RecordBatch(num_rows=3, schema=Schema(fields=[dept_id: int64, dept: string]))

Step 1 — filter rows

Filtering rows has two spellings, and it is worth seeing both.

The query layer is the one you will reach for. col(...) names a column, comparison operators build a predicate, and the engine applies it to every column for you:

from marrow import col, lit

top_earners = (
    ma.memtable(employees)
      .filter(col("salary") > lit(80_000))
      .collect()
)
print(top_earners.to_pylist())
[{'name': 'Alice', 'dept_id': 1, 'salary': 95000, 'age': 34}, {'name': 'Carol', 'dept_id': 1, 'salary': 120000, 'age': 45}, {'name': 'Dave', 'dept_id': 3, 'salary': 88000, 'age': 52}, {'name': 'Frank', 'dept_id': 1, 'salary': 105000, 'age': 39}]

The eager kernels work a column at a time, so the same thing by hand is a mask plus a filter per column. Reach for this when you already hold a mask:

threshold = ma.array([80_000] * employees.num_rows)
well_paid = mc.greater(employees.column("salary"), threshold)
print("mask:", well_paid)

kept = mc.filter(employees.column("name"), well_paid)
print("names:", kept)
mask: BoolArray([True, False, True, True, False, True])
names: StringArray([Alice, Carol, Dave, Frank])

Note the comparison needs both operands to share a dtype — hence the array of thresholds. The query layer handles that for you, which is most of why it is the shorter spelling.

Step 2 — sort

sort_by orders a batch by one or more keys. Pass a list of (column, "ascending" | "descending") tuples; the second argument controls null placement (None uses the default). Multi-column keys are honoured — ties on the first key break on the next:

by_salary = top_earners.sort_by([("dept_id", "ascending"), ("salary", "descending")])
print(by_salary.to_pylist())
[{'name': 'Carol', 'dept_id': 1, 'salary': 120000, 'age': 45}, {'name': 'Frank', 'dept_id': 1, 'salary': 105000, 'age': 39}, {'name': 'Alice', 'dept_id': 1, 'salary': 95000, 'age': 34}, {'name': 'Dave', 'dept_id': 3, 'salary': 88000, 'age': 52}]

Step 3 — join

join performs a hash join between two record batches on named key columns. The signature is join(right, keys, right_keys=None, join_type="inner", num_threads=0):

  • keys — list of left-side key column names,
  • right_keys — right-side names, or None to reuse keys,
  • join_type — "inner", "left", "right", "full", "semi" or "anti",
  • num_threads — 0 auto-selects cores, 1 forces the serial path.

Enrich each employee with their department name:

enriched = employees.join(departments, ["dept_id"], join_type="inner")
print(enriched.to_pylist())
[{'name': 'Alice', 'dept_id': 1, 'salary': 95000, 'age': 34, 'dept_id_right': 1, 'dept': 'Engineering'}, {'name': 'Carol', 'dept_id': 1, 'salary': 120000, 'age': 45, 'dept_id_right': 1, 'dept': 'Engineering'}, {'name': 'Frank', 'dept_id': 1, 'salary': 105000, 'age': 39, 'dept_id_right': 1, 'dept': 'Engineering'}, {'name': 'Bob', 'dept_id': 2, 'salary': 72000, 'age': 28, 'dept_id_right': 2, 'dept': 'Sales'}, {'name': 'Eve', 'dept_id': 2, 'salary': 61000, 'age': 24, 'dept_id_right': 2, 'dept': 'Sales'}, {'name': 'Dave', 'dept_id': 3, 'salary': 88000, 'age': 52, 'dept_id_right': 3, 'dept': 'Marketing'}]

Columns that collide between the two sides are suffixed with _right, the right-hand key included.

Step 4 — put it together

The whole pipeline reads as a short sequence of transformations: join, filter, sort, project.

# 1. join employees to their departments
joined = employees.join(departments, ["dept_id"], join_type="inner")

# 2-4. filter, sort and project as one plan
report = (
    ma.memtable(joined)
      .filter(col("salary") > lit(80_000))
      .order_by(("salary", "descending"))
      .select("name", "dept", "salary")
      .collect()
)
print(report.to_pylist())
[{'name': 'Carol', 'dept': 'Engineering', 'salary': 120000}, {'name': 'Frank', 'dept': 'Engineering', 'salary': 105000}, {'name': 'Alice', 'dept': 'Engineering', 'salary': 95000}, {'name': 'Dave', 'dept': 'Marketing', 'salary': 88000}]

Only the join is eager here — it is a RecordBatch method. Everything after it is one plan, which runs when collect() is called. Add .optimize() before collect() to let the plan rewriter act on it first.

Step 5 — aggregate

Aggregation is not an eager kernel call — it happens inside a query plan. Wrap the batch with ma.memtable() to get a lazy table, describe what you want, and call .collect() to run it. Nulls are skipped:

q = ma.memtable(joined)

totals = q.aggregate(
    headcount=("count", "salary"),
    total_spend=("sum", "salary"),
    top_salary=("max", "salary"),
)
print(totals.collect().to_pylist())
[{'headcount': 6, 'total_spend': 541000, 'top_salary': 120000}]

The same call with a by= key groups instead of reducing to one row — one code path for both:

by_dept = (
    q.aggregate(
        by=["dept"],
        headcount=("count", "salary"),
        avg_salary=("mean", "salary"),
    )
    .order_by("dept")
)
print(by_dept.collect().to_pylist())
[{'dept': 'Engineering', 'headcount': 3, 'avg_salary': 106666.66666666667}, {'dept': 'Marketing', 'headcount': 1, 'avg_salary': 88000.0}, {'dept': 'Sales', 'headcount': 2, 'avg_salary': 66500.0}]

Nothing has run until .collect(), so a plan is a reusable template — and explain() shows what will happen:

print(by_dept.explain())
Sort(Aggregate(InMemoryTable(6 rows), by=dept, count(salary), mean(salary)), dept asc)

Where to go next

Back to top