Run a remote query

Polars Nanook runs an existing Polars LazyFrame remotely without changes to the query: add .remote(context=ctx) to it and call .collect(). The call after .remote() decides where the result goes:

.collect() returns a DataFrame and .collect_batches() an iterator of DataFrames, both streamed back to the client.
.sink_parquet() and the other sinks write the result to storage.
.execute() stores the result at the cluster's results location for later use, as described in Keep results on the cluster.

The local query

The following example uses a query from the PDS-H benchmark suite, which is derived from the TPC-H benchmark. It runs with the open-source Polars library on a local machine:

import polars as pl

customer = pl.scan_parquet("data/customer.parquet")
lineitem = pl.scan_parquet("data/lineitem.parquet")
orders = pl.scan_parquet("data/orders.parquet")


def pdsh_q3(
    customer: pl.LazyFrame, lineitem: pl.LazyFrame, orders: pl.LazyFrame
) -> pl.LazyFrame:
    return (
        customer.filter(pl.col("c_mktsegment") == "BUILDING")
        .join(orders, left_on="c_custkey", right_on="o_custkey")
        .join(lineitem, left_on="o_orderkey", right_on="l_orderkey")
        .filter(pl.col("o_orderdate") < pl.date(1995, 3, 15))
        .filter(pl.col("l_shipdate") > pl.date(1995, 3, 15))
        .with_columns(
            (pl.col("l_extendedprice") * (1 - pl.col("l_discount"))).alias("revenue")
        )
        .group_by("o_orderkey", "o_orderdate", "o_shippriority")
        .agg(pl.sum("revenue"))
        .select(
            pl.col("o_orderkey").alias("l_orderkey"),
            "revenue",
            "o_orderdate",
            "o_shippriority",
        )
        .sort(by=["revenue", "o_orderdate"], descending=[True, False])
    )


pdsh_q3(customer, lineitem, orders).collect()

Running the query on Polars Nanook

The same query runs on a larger dataset in S3 and returns the result as a DataFrame:

import polars_nanook as pn

ctx: pn.ComputeContext | pn.ClusterContext

# Use a larger dataset available on S3
lineitem_sf10 = pl.scan_parquet(
    "s3://polars-cloud-samples-us-east-2-prd/pdsh/sf10/lineitem.parquet",
    storage_options={"request_payer": "true", "aws_region": "us-east-2"},
)
customer_sf10 = pl.scan_parquet(
    "s3://polars-cloud-samples-us-east-2-prd/pdsh/sf10/customer.parquet",
    storage_options={"request_payer": "true", "aws_region": "us-east-2"},
)
orders_sf10 = pl.scan_parquet(
    "s3://polars-cloud-samples-us-east-2-prd/pdsh/sf10/orders.parquet",
    storage_options={"request_payer": "true", "aws_region": "us-east-2"},
)

# Your query remains the same
pdsh_q3(customer_sf10, lineitem_sf10, orders_sf10).remote(context=ctx).collect()

Run the examples yourself

The examples read the public datasets in us-east-2, a Requester Pays bucket. That page covers the costs, the credentials the workers need and how the region affects latency.

SQL queries

pl.sql and SQLContext.execute return a LazyFrame, so .remote() works on them as on any other query.

sql

import polars_nanook as pn

ctx: pn.ComputeContext | pn.ClusterContext

query = pl.sql(
    """
    SELECT
        l_orderkey,
        SUM(l_extendedprice * (1 - l_discount)) AS revenue,
        o_orderdate,
        o_shippriority
    FROM customer_sf10
    JOIN orders_sf10 ON c_custkey = o_custkey
    JOIN lineitem_sf10 ON l_orderkey = o_orderkey
    WHERE c_mktsegment = 'BUILDING'
        AND o_orderdate < DATE '1995-03-15'
        AND l_shipdate > DATE '1995-03-15'
    GROUP BY l_orderkey, o_orderdate, o_shippriority
    ORDER BY revenue DESC, o_orderdate
    """
)

query.remote(context=ctx).collect()

Writing results to storage

For batch processing, .sink_parquet() writes results directly to storage reachable from the workers, for example an S3 bucket. A sink returns as soon as the query is submitted. .await_result() then waits for the query to finish and raises if it fails:

import polars_nanook as pn

ctx: pn.ComputeContext | pn.ClusterContext

# Replace the S3 url with your own to run the query successfully
(
    pdsh_q3(customer_sf10, lineitem_sf10, orders_sf10)
    .remote(context=ctx)
    .sink_parquet("s3://<BUCKET>/processed-data/")
    .await_result()
)

A path without a file extension, as in the example, is a directory in which every worker writes its own files. A path with one, such as out.parquet, needs sink_to_single_file=True in the sink call and is written by one worker. .collect() also gathers the result on one worker, so sink large results to a directory.