From 9420865bbe1776647bebda3a81a86325854f6a7c Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Mon, 23 Dec 2024 12:59:16 +0100 Subject: [PATCH] Add option to run Polars on the new-streaming engine. (#138) --- queries/polars/utils.py | 9 +++++---- settings.py | 1 + 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/queries/polars/utils.py b/queries/polars/utils.py index 4dc3916..200dfef 100644 --- a/queries/polars/utils.py +++ b/queries/polars/utils.py @@ -127,20 +127,21 @@ def obtain_engine_config() -> pl.GPUEngine | Literal["cpu"]: def run_query(query_number: int, lf: pl.LazyFrame) -> None: streaming = settings.run.polars_streaming + new_streaming = settings.run.polars_new_streaming eager = settings.run.polars_eager gpu = settings.run.polars_gpu - if (eager or streaming) and gpu: - msg = "polars-gpu engine does not support eager or streaming" + if sum([eager, streaming, new_streaming, gpu]) > 1: + msg = "Please specify at most one of eager, streaming, new_streaming or gpu" raise ValueError(msg) if settings.run.polars_show_plan: - print(lf.explain(streaming=streaming, optimized=eager)) + print(lf.explain(streaming=streaming, new_streaming=new_streaming, optimized=eager)) engine = obtain_engine_config() # Eager load engine backend, so we don't time that. _preload_engine(engine) query = partial( - lf.collect, streaming=streaming, no_optimization=eager, engine=engine + lf.collect, streaming=streaming, new_streaming=new_streaming, no_optimization=eager, engine=engine ) if gpu: diff --git a/settings.py b/settings.py index c38d54e..c6e142e 100644 --- a/settings.py +++ b/settings.py @@ -31,6 +31,7 @@ class Run(BaseSettings): polars_show_plan: bool = False polars_eager: bool = False polars_streaming: bool = False + polars_new_streaming: bool = False polars_gpu: bool = False # Use GPU engine? polars_gpu_device: int = 0 # The GPU device to run on for polars GPU # Which style of GPU memory resource to use