corvic-engine 0.3.0rc67__cp38-abi3-win_amd64.whl → 0.3.0rc68__cp38-abi3-win_amd64.whl

corvic/system/in_memory_executor.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
+import asyncio
 import dataclasses
 import datetime
 import functools
-import math
-from collections.abc import Callable, Mapping, MutableMapping
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import AbstractContextManager, ExitStack, nullcontext
 from types import TracebackType
-from typing import Any, Final, cast
+from typing import Any, Final, ParamSpec, Self, TypeVar, cast
 
 import numpy as np
 import polars as pl
@@ -19,8 +20,9 @@ import pyarrow.parquet as pq
 import structlog
 from google.protobuf import json_format, struct_pb2
 from more_itertools import flatten
-from typing_extensions import Self, deprecated
+from typing_extensions import deprecated
 
+import corvic.system._column_encoding as column_encoding
 from corvic import embed, embedding_metric, op_graph, sql
 from corvic.result import (
     InternalError,
@@ -49,48 +51,51 @@ from corvic_generated.orm.v1 import table_pb2
 
 _logger = structlog.get_logger()
 
-"""Reference and Maximum number of years for normalizing year in Datetime encoder"""
-REFERENCE_YEAR: Final = 1900
-MAX_NUMBER_OF_YEARS: Final = 200
-
 _MIN_EMBEDDINGS_FOR_EMBEDDINGS_SUMMARY: Final = 3
 
 
+def _collect_and_apply(
+    data: pl.LazyFrame, func: Callable[[pl.DataFrame], _R]
+) -> tuple[pl.DataFrame, _R]:
+    data_df = data.collect()
+    return data_df, func(data_df)
+
+
 def get_polars_embedding_length(
-    embedding_df: pl.DataFrame, embedding_column_name: str
+    embedding_series: pl.Series,
 ) -> Ok[int] | InvalidArgumentError:
-    outer_type = embedding_df.schema[embedding_column_name]
+    outer_type = embedding_series.dtype
     if isinstance(outer_type, pl.Array):
         return Ok(outer_type.shape[0])
     if not isinstance(outer_type, pl.List):
         return InvalidArgumentError("invalid embedding datatype", dtype=str(outer_type))
-    if len(embedding_df[embedding_column_name]) == 0:
+    if len(embedding_series) == 0:
         return InvalidArgumentError(
             "cannot infer embedding length for empty embedding set"
         )
-    embedding_length = len(embedding_df[embedding_column_name][0])
+    embedding_length = len(embedding_series[0])
     if embedding_length < 1:
         return InvalidArgumentError("invalid embedding length", length=embedding_length)
     return Ok(embedding_length)
 
 
 def get_polars_embedding(
-    embedding_df: pl.DataFrame, embedding_column_name: str
+    embedding_series: pl.Series,
 ) -> Ok[np.ndarray[Any, Any]] | InvalidArgumentError:
-    outer_type = embedding_df.schema[embedding_column_name]
+    outer_type = embedding_series.dtype
     if isinstance(outer_type, pl.Array):
-        return Ok(embedding_df[embedding_column_name].to_numpy())
+        return Ok(embedding_series.to_numpy())
     if not isinstance(outer_type, pl.List):
         return InvalidArgumentError("invalid embedding datatype", dtype=str(outer_type))
-    match get_polars_embedding_length(embedding_df, embedding_column_name):
+    match get_polars_embedding_length(embedding_series):
         case Ok(embedding_length):
             pass
         case InvalidArgumentError() as err:
             return err
     return Ok(
-        embedding_df[embedding_column_name]
-        .cast(pl.Array(inner=outer_type.inner, shape=embedding_length))
-        .to_numpy()
+        embedding_series.cast(
+            pl.Array(inner=outer_type.inner, shape=embedding_length)
+        ).to_numpy()
     )
 
 
@@ -192,8 +197,15 @@ class _SchemaAndBatches:
     metrics: dict[str, Any]
 
     @classmethod
-    def from_lazy_frame_with_metrics(cls, lfm: _LazyFrameWithMetrics):
-        return cls.from_dataframe(lfm.data.collect(), lfm.metrics)
+    async def from_lazy_frame_with_metrics(
+        cls, lfm: _LazyFrameWithMetrics, worker_threads: ThreadPoolExecutor | None
+    ):
+        return cls.from_dataframe(
+            await asyncio.get_running_loop().run_in_executor(
+                worker_threads, lfm.data.collect
+            ),
+            lfm.metrics,
+        )
 
     def to_batch_reader(self):
         return pa.RecordBatchReader.from_batches(
@@ -221,6 +233,10 @@ class _SchemaAndBatches:
         return cls(schema, table.to_batches(), metrics)
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
 @dataclasses.dataclass(frozen=True)
 class _SlicedTable:
     op_graph: op_graph.Op
@@ -230,6 +246,7 @@ class _SlicedTable:
 @dataclasses.dataclass
 class _InMemoryExecutionContext(AbstractContextManager["_InMemoryExecutionContext"]):
     exec_context: ExecutionContext
+    worker_threads: ThreadPoolExecutor | None
     current_output_context: TableComputeContext | None = None
 
     # Using _SchemaAndBatches rather than a RecordBatchReader since the latter's
@@ -239,6 +256,7 @@ class _InMemoryExecutionContext(AbstractContextManager["_InMemoryExecutionContex
         dataclasses.field(default_factory=dict)
     )
     exit_stack: ExitStack = dataclasses.field(default_factory=ExitStack)
+    lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock)
 
     def __enter__(self) -> Self:
         self.exit_stack = self.exit_stack.__enter__()
@@ -252,6 +270,15 @@ class _InMemoryExecutionContext(AbstractContextManager["_InMemoryExecutionContex
     ) -> bool | None:
         return self.exit_stack.__exit__(__exc_type, __exc_value, __traceback)
 
+    async def run_on_worker(
+        self, func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs
+    ) -> _R:
+        # lock here because polars operations aren't guaranteed to be independent
+        async with self.lock:
+            return await asyncio.get_running_loop().run_in_executor(
+                self.worker_threads, lambda: func(*args, **kwargs)
+            )
+
     @classmethod
     def count_source_op_uses(
         cls,
@@ -407,55 +434,55 @@ class InMemoryExecutor(OpGraphExecutor):
         )
         return Ok(_LazyFrameWithMetrics(data, metrics={}))
 
-    def _execute_rollup_by_aggregation(
+    async def _execute_rollup_by_aggregation(
         self, op: op_graph.op.RollupByAggregation, context: _InMemoryExecutionContext
     ) -> Ok[_LazyFrameWithMetrics]:
         raise NotImplementedError(
             "rollup by aggregation outside of sql not implemented"
         )
 
-    def _compute_source_then_apply(
+    async def _compute_source_then_apply(
         self,
         source: op_graph.Op,
         lf_op: Callable[[pl.LazyFrame], pl.LazyFrame],
         context: _InMemoryExecutionContext,
     ):
-        return self._execute(source, context).map(
+        return (await self._execute(source, context)).map(
             lambda source_lfm: source_lfm.apply(lf_op)
         )
 
-    def _execute_rename_columns(
+    async def _execute_rename_columns(
         self, op: op_graph.op.RenameColumns, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source, lambda lf: lf.rename(dict(op.old_name_to_new)), context
         )
 
-    def _execute_select_columns(
+    async def _execute_select_columns(
         self, op: op_graph.op.SelectColumns, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source, lambda lf: lf.select(op.columns), context
         )
 
-    def _execute_limit_rows(
+    async def _execute_limit_rows(
         self, op: op_graph.op.LimitRows, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source, lambda lf: lf.limit(op.num_rows), context
         )
 
-    def _execute_offset_rows(
+    async def _execute_offset_rows(
         self, op: op_graph.op.OffsetRows, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source, lambda lf: lf.slice(op.num_rows), context
         )
 
-    def _execute_order_by(
+    async def _execute_order_by(
         self, op: op_graph.op.OrderBy, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source, lambda lf: lf.sort(op.columns, descending=op.desc), context
         )
 
@@ -536,7 +563,7 @@ class InMemoryExecutor(OpGraphExecutor):
             case op_graph.row_filter.CombineFilters():
                 return self._row_filter_combination_to_condition(row_filter)
 
-    def _execute_filter_rows(
+    async def _execute_filter_rows(
         self, op: op_graph.op.FilterRows, context: _InMemoryExecutionContext
     ):
         match self._row_filter_to_condition(op.row_filter):
@@ -544,78 +571,129 @@ class InMemoryExecutor(OpGraphExecutor):
                 return InternalError.from_(err)
             case Ok(row_filter):
                 pass
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source, lambda lf: lf.filter(row_filter), context
         )
 
-    def _execute_embedding_metrics(  # noqa: C901
+    def _get_embedding_column_if(
+        self,
+        data: pl.LazyFrame,
+        embedding_column_name: str,
+        pred: Callable[[pl.DataFrame], bool],
+    ) -> Ok[tuple[pl.DataFrame, np.ndarray[Any, Any] | None]] | InvalidArgumentError:
+        data_df = data.collect()
+        if pred(data_df):
+            match get_polars_embedding(data_df[embedding_column_name]):
+                case InvalidArgumentError() as err:
+                    return err
+                case Ok(embeddings):
+                    pass
+        else:
+            embeddings = None
+        return Ok((data_df, embeddings))
+
+    async def _execute_embedding_metrics(  # noqa: C901
         self, op: op_graph.op.EmbeddingMetrics, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.table, context):
-            case Ok(source_lfm):
-                pass
-            case err:
-                return err
-        embedding_df = source_lfm.data.collect()
-
-        if len(embedding_df) < _MIN_EMBEDDINGS_FOR_EMBEDDINGS_SUMMARY:
-            # downstream consumers handle empty metadata by substituting their
-            # own values
-            return Ok(
-                _LazyFrameWithMetrics(embedding_df.lazy(), metrics=source_lfm.metrics)
-            )
-
         # before it was configurable, this op assumed that the column's name was
         # this hardcoded name
         embedding_column_name = op.embedding_column_name or "embedding"
-        match get_polars_embedding(embedding_df, embedding_column_name):
-            case Ok(embedding):
+
+        match await self._execute(op.table, context):
+            case Ok(source_lfm):
                 pass
+            case err:
+                return err
+        match await context.run_on_worker(
+            self._get_embedding_column_if,
+            source_lfm.data,
+            embedding_column_name,
+            lambda df: len(df) >= _MIN_EMBEDDINGS_FOR_EMBEDDINGS_SUMMARY,
+        ):
             case InvalidArgumentError() as err:
-                return InternalError.from_(err)
+                return err
+            case Ok(result):
+                embedding_df, embeddings = result
+
+        if embeddings is None:
+            return Ok(_LazyFrameWithMetrics(embedding_df.lazy(), source_lfm.metrics))
 
         metrics = source_lfm.metrics.copy()
-        match embedding_metric.ne_sum(embedding, normalize=True):
+        async with asyncio.TaskGroup() as tg:
+            ne_sum = tg.create_task(
+                context.run_on_worker(
+                    embedding_metric.ne_sum, embeddings, normalize=True
+                )
+            )
+            condition_number = tg.create_task(
+                context.run_on_worker(
+                    embedding_metric.condition_number, embeddings, normalize=True
+                )
+            )
+            rcondition_number = tg.create_task(
+                context.run_on_worker(
+                    embedding_metric.rcondition_number, embeddings, normalize=True
+                )
+            )
+            stable_rank = tg.create_task(
+                context.run_on_worker(
+                    embedding_metric.stable_rank, embeddings, normalize=True
+                )
+            )
+
+        match ne_sum.result():
             case Ok(metric):
                 metrics["ne_sum"] = metric
             case InvalidArgumentError() as err:
                 _logger.warning("could not compute ne_sum", exc_info=str(err))
-        match embedding_metric.condition_number(embedding, normalize=True):
+
+        match condition_number.result():
             case Ok(metric):
                 metrics["condition_number"] = metric
             case InvalidArgumentError() as err:
                 _logger.warning("could not compute condition_number", exc_info=str(err))
-        match embedding_metric.rcondition_number(embedding, normalize=True):
+
+        match rcondition_number.result():
             case Ok(metric):
                 metrics["rcondition_number"] = metric
             case InvalidArgumentError() as err:
                 _logger.warning(
                     "could not compute rcondition_number", exc_info=str(err)
                 )
-        match embedding_metric.stable_rank(embedding, normalize=True):
+        match stable_rank.result():
             case Ok(metric):
                 metrics["stable_rank"] = metric
             case InvalidArgumentError() as err:
                 _logger.warning("could not compute stable_rank", exc_info=str(err))
+
         return Ok(_LazyFrameWithMetrics(embedding_df.lazy(), metrics=metrics))
 
-    def _execute_embedding_coordinates(
+    async def _execute_embedding_coordinates(
         self, op: op_graph.op.EmbeddingCoordinates, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.table, context):
+        match await self._execute(op.table, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        embedding_df = source_lfm.data.collect()
 
         # before it was configurable, this op assumed that the column's name was
         # this hardcoded name
         embedding_column_name = op.embedding_column_name or "embedding"
+        match await context.run_on_worker(
+            self._get_embedding_column_if,
+            source_lfm.data,
+            embedding_column_name,
+            lambda df: len(df) >= _MIN_EMBEDDINGS_FOR_EMBEDDINGS_SUMMARY,
+        ):
+            case InvalidArgumentError() as err:
+                return err
+            case Ok(result):
+                embedding_df, embeddings = result
 
         # the neighbors of a point includes itself. That does mean, that an n_neighbors
         # value of less than 3 simply does not work
-        if len(embedding_df) < _MIN_EMBEDDINGS_FOR_EMBEDDINGS_SUMMARY:
+        if embeddings is None:
             coordinates_df = embedding_df.lazy().with_columns(
                 pl.Series(
                     name=embedding_column_name,
@@ -625,14 +703,11 @@ class InMemoryExecutor(OpGraphExecutor):
             )
             return Ok(_LazyFrameWithMetrics(coordinates_df, source_lfm.metrics))
 
-        match get_polars_embedding(embedding_df, embedding_column_name):
-            case Ok(embedding):
-                pass
-            case InvalidArgumentError() as err:
-                raise err
-
-        match self._dimension_reducer.reduce_dimensions(
-            embedding, op.n_components, op.metric
+        match await context.run_on_worker(
+            self._dimension_reducer.reduce_dimensions,
+            embeddings,
+            op.n_components,
+            op.metric,
         ):
             case Ok(coordinates):
                 pass
@@ -648,26 +723,37 @@ class InMemoryExecutor(OpGraphExecutor):
         )
         return Ok(_LazyFrameWithMetrics(coordinates_df, source_lfm.metrics))
 
-    def _execute_distinct_rows(
+    async def _execute_distinct_rows(
         self, op: op_graph.op.DistinctRows, context: _InMemoryExecutionContext
     ):
-        return self._execute(op.source, context).map(
-            lambda source_lfm: _LazyFrameWithMetrics(
-                source_lfm.data.unique(), source_lfm.metrics
-            )
+        return await self._compute_source_then_apply(
+            op.source, lambda source: source.unique(), context
         )
 
-    def _execute_join(self, op: op_graph.op.Join, context: _InMemoryExecutionContext):
-        match self._execute(op.left_source, context):
+    async def _execute_join(
+        self, op: op_graph.op.Join, context: _InMemoryExecutionContext
+    ):
+        async with asyncio.TaskGroup() as tg:
+            left_task = tg.create_task(self._execute(op.left_source, context))
+            right_task = tg.create_task(self._execute(op.right_source, context))
+        match left_task.result():
+            case (
+                InternalError()
+                | ResourceExhaustedError()
+                | InvalidArgumentError() as err
+            ):
+                return err
             case Ok(left_lfm):
                 pass
-            case err:
+        match right_task.result():
+            case (
+                InternalError()
+                | ResourceExhaustedError()
+                | InvalidArgumentError() as err
+            ):
                 return err
-        match self._execute(op.right_source, context):
             case Ok(right_lfm):
                 pass
-            case err:
-                return err
         left_lf = left_lfm.data
         right_lf = right_lfm.data
 
@@ -702,20 +788,31 @@ class InMemoryExecutor(OpGraphExecutor):
             )
         )
 
-    def _execute_empty(self, op: op_graph.op.Empty, context: _InMemoryExecutionContext):
+    async def _execute_empty(
+        self, op: op_graph.op.Empty, context: _InMemoryExecutionContext
+    ):
         empty_table = cast(pl.DataFrame, pl.from_arrow(pa.schema([]).empty_table()))
         return Ok(_LazyFrameWithMetrics(empty_table.lazy(), metrics={}))
 
-    def _execute_concat(
+    async def _execute_concat(
         self, op: op_graph.op.Concat, context: _InMemoryExecutionContext
     ):
+        async with asyncio.TaskGroup() as tg:
+            tasks = [
+                tg.create_task(self._execute(table, context)) for table in op.tables
+            ]
+
         source_lfms = list[_LazyFrameWithMetrics]()
-        for table in op.tables:
-            match self._execute(table, context):
-                case Ok(batches):
-                    source_lfms.append(batches)
-                case err:
+        for task in tasks:
+            match task.result():
+                case (
+                    InternalError()
+                    | ResourceExhaustedError()
+                    | InvalidArgumentError() as err
+                ):
                     return err
+                case Ok(lfm):
+                    source_lfms.append(lfm)
         data = pl.concat([lfm.data for lfm in source_lfms], how=op.how)
         metrics = dict[str, Any]()
         for lfm in source_lfms:
@@ -794,16 +891,19 @@ class InMemoryExecutor(OpGraphExecutor):
             context,
         )
 
-    def _execute_embed_column(
+    async def _execute_embed_column(
         self, op: op_graph.op.EmbedColumn, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
-        to_embed = source_df[op.column_name].cast(pl.String())
+        source_df, to_embed = await context.run_on_worker(
+            _collect_and_apply,
+            source_lfm.data,
+            lambda df: df[op.column_name].cast(pl.String),
+        )
 
         embed_context = EmbedTextContext(
             inputs=to_embed,
@@ -813,112 +913,19 @@ class InMemoryExecutor(OpGraphExecutor):
             expected_coordinate_bitwidth=op.expected_coordinate_bitwidth,
             room_id=context.exec_context.room_id,
         )
-        match self._text_embedder.embed(embed_context):
+        match await self._text_embedder.aembed(embed_context, context.worker_threads):
             case Ok(result):
                 pass
             case InvalidArgumentError() | InternalError() as err:
                 raise InternalError("Failed to embed column") from err
 
-        result_df = (
-            source_df.lazy()
-            .with_columns(result.embeddings.alias(op.embedding_column_name))
-            .drop_nulls(op.embedding_column_name)
-        )
-
-        return Ok(source_lfm.with_data(result_df))
-
-    @staticmethod
-    def get_cyclic_encoding(
-        series: pl.Series,
-        period: int,
-    ) -> tuple[pl.Series, pl.Series]:
-        sine_series = (2 * math.pi * series / period).sin().alias(f"{series.name}_sine")
-        cosine_series = (
-            (2 * math.pi * series / period).cos().alias(f"{series.name}_cosine")
+        return Ok(
+            source_lfm.with_data(
+                source_df.lazy()
+                .with_columns(result.embeddings.alias(op.embedding_column_name))
+                .drop_nulls(op.embedding_column_name)
+            )
         )
-        return sine_series, cosine_series
-
-    @staticmethod
-    def encode_datetime(series: pl.Series) -> pl.Series:
-        match series.dtype:
-            case pl.Date | pl.Time:
-                pass
-            case pl.Datetime:
-                series = series.dt.replace_time_zone("UTC")
-            case _:
-                raise ValueError("Invalid arguments, expected a datetime series")
-
-        if series.is_null().all():
-            zero_vector = pl.zeros(11, dtype=pl.Float32, eager=True)
-            return pl.Series([zero_vector] * len(series), dtype=pl.List(pl.Float32))
-
-        n = len(series)
-        year_norm = pl.zeros(n, dtype=pl.Float32, eager=True).alias("year")
-        month_sine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("month_sine")
-        month_cosine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("month_cosine")
-        day_sine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("day_sine")
-        day_cosine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("day_cosine")
-        hour_sine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("hour_sine")
-        hour_cosine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("hour_cosine")
-        minute_sine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("minute_sine")
-        minute_cosine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("minute_cosine")
-        second_sine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("second_sine")
-        second_cosine = pl.zeros(n, dtype=pl.Float32, eager=True).alias("second_cosine")
-
-        if series.dtype in [pl.Date, pl.Datetime]:
-            try:
-                year = series.dt.year().cast(pl.Float32).alias("year")
-                month = series.dt.month().cast(pl.Float32).alias("month")
-                day = series.dt.day().cast(pl.Float32).alias("day")
-
-                year_norm = (year - REFERENCE_YEAR) / MAX_NUMBER_OF_YEARS
-                month_sine, month_cosine = InMemoryExecutor.get_cyclic_encoding(
-                    month, 12
-                )
-                day_sine, day_cosine = InMemoryExecutor.get_cyclic_encoding(day, 31)
-            except pl.exceptions.PanicException as e:
-                _logger.exception("Error extracting datetime", exc_info=e)
-
-        if series.dtype in [pl.Time, pl.Datetime]:
-            try:
-                hour = series.dt.hour().cast(pl.Float32).alias("hour")
-                minute = series.dt.minute().cast(pl.Float32).alias("minute")
-                second = series.dt.second().cast(pl.Float32).alias("second")
-
-                hour_sine, hour_cosine = InMemoryExecutor.get_cyclic_encoding(hour, 24)
-                minute_sine, minute_cosine = InMemoryExecutor.get_cyclic_encoding(
-                    minute, 60
-                )
-                second_sine, second_cosine = InMemoryExecutor.get_cyclic_encoding(
-                    second, 60
-                )
-            except pl.exceptions.PanicException as e:
-                _logger.exception("Error extracting datetime", exc_info=e)
-
-        return pl.DataFrame(
-            [
-                year_norm.fill_null(0.0),
-                month_sine.fill_null(0.0),
-                month_cosine.fill_null(0.0),
-                day_sine.fill_null(0.0),
-                day_cosine.fill_null(0.0),
-                hour_sine.fill_null(0.0),
-                hour_cosine.fill_null(0.0),
-                minute_sine.fill_null(0.0),
-                minute_cosine.fill_null(0.0),
-                second_sine.fill_null(0.0),
-                second_cosine.fill_null(0.0),
-            ]
-        ).select(pl.concat_list(pl.all()).alias(series.name))[series.name]
-
-    @staticmethod
-    def encode_duration(series: pl.Series) -> pl.Series:
-        if series.dtype != pl.Duration:
-            raise ValueError("Invalid arguments, expected a duration series")
-        if series.is_null().all():
-            return pl.zeros(len(series), dtype=pl.Float32, eager=True)
-
-        return series.dt.total_seconds().cast(pl.Float32).fill_null(0.0)
 
     @staticmethod
     def encode_text(series: pl.Series) -> pl.Series:
@@ -939,132 +946,71 @@ class InMemoryExecutor(OpGraphExecutor):
             pl.List(pl.Float32),
         )
 
-    def _execute_encode_columns(  # noqa: C901, PLR0915
+    def _encode_column(  # noqa: C901
+        self, to_encode: pl.Series, encoder: op_graph.Encoder
+    ) -> tuple[pl.Series, list[str] | None]:
+        match encoder:
+            case op_graph.encoder.OneHotEncoder():
+                return column_encoding.encode_one_hot(to_encode)
+
+            case op_graph.encoder.MinMaxScaler():
+                return column_encoding.encode_min_max_scale(
+                    to_encode, encoder.feature_range_min, encoder.feature_range_max
+                ), None
+
+            case op_graph.encoder.LabelBinarizer():
+                return column_encoding.encode_label_boolean(
+                    to_encode, encoder.neg_label, encoder.pos_label
+                ), None
+
+            case op_graph.encoder.LabelEncoder():
+                return column_encoding.encode_label(
+                    to_encode, normalize=encoder.normalize
+                ), None
+
+            case op_graph.encoder.KBinsDiscretizer():
+                return column_encoding.encode_kbins(
+                    to_encode, encoder.n_bins, encoder.encode_method, encoder.strategy
+                ), None
+
+            case op_graph.encoder.Binarizer():
+                return column_encoding.encode_boolean(
+                    to_encode, encoder.threshold
+                ), None
+
+            case op_graph.encoder.MaxAbsScaler():
+                return column_encoding.encode_max_abs_scale(to_encode), None
+
+            case op_graph.encoder.StandardScaler():
+                return column_encoding.encode_standard_scale(
+                    to_encode, with_mean=encoder.with_mean, with_std=encoder.with_std
+                ), None
+
+            case op_graph.encoder.TimestampEncoder():
+                if to_encode.dtype == pl.datatypes.Duration:
+                    return column_encoding.encode_duration(to_encode), None
+                return column_encoding.encode_datetime(to_encode), None
+
+            case op_graph.encoder.TextEncoder():
+                return self.encode_text(to_encode), None
+
+    async def _execute_encode_columns(
         self, op: op_graph.op.EncodeColumns, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
+        source_df = await context.run_on_worker(source_lfm.data.collect)
         metrics = source_lfm.metrics.copy()
         metric = metrics.get("one_hot_encoder", {})
         for encoder_arg in op.encoded_columns:
-            to_encode = source_df[encoder_arg.column_name]
-            match encoder_arg.encoder:
-                case op_graph.encoder.OneHotEncoder():
-                    encoded = to_encode.to_dummies()
-                    metric[encoder_arg.column_name] = encoded.columns
-                    encoded = encoded.select(
-                        pl.concat_list(pl.all())
-                        .alias(encoder_arg.encoded_column_name)
-                        .cast(pl.List(pl.Boolean))
-                    )
-
-                case op_graph.encoder.MinMaxScaler():
-                    from sklearn.preprocessing import MinMaxScaler
-
-                    encoder = MinMaxScaler(
-                        feature_range=(
-                            encoder_arg.encoder.feature_range_min,
-                            encoder_arg.encoder.feature_range_max,
-                        )
-                    )
-                    encoded = encoder.fit_transform(
-                        to_encode.to_numpy().reshape(-1, 1)
-                    ).flatten()
-
-                case op_graph.encoder.LabelBinarizer():
-                    from sklearn.preprocessing import LabelBinarizer
-
-                    encoder = LabelBinarizer(
-                        neg_label=encoder_arg.encoder.neg_label,
-                        pos_label=encoder_arg.encoder.pos_label,
-                    )
-                    encoded = encoder.fit_transform(to_encode.to_numpy().reshape(-1))
-
-                case op_graph.encoder.LabelEncoder():
-                    from sklearn.preprocessing import LabelEncoder
-
-                    encoder = LabelEncoder()
-                    encoded = encoder.fit_transform(
-                        to_encode.to_numpy().reshape(-1)
-                    ).flatten()
-                    # `classes_` is only set after fit,
-                    # Creating custom typestubs will not solve this typing issue.
-                    if encoder_arg.encoder.normalize and hasattr(encoder, "classes_"):
-                        classes_ = cast(list[int], encoder.classes_)  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
-                        max_class: int = len(classes_) - 1
-                        if max_class > 0:
-                            encoded = encoded.astype(np.float64)
-                            encoded /= max_class
-
-                case op_graph.encoder.KBinsDiscretizer():
-                    from sklearn.preprocessing import KBinsDiscretizer
-
-                    encoder = KBinsDiscretizer(
-                        n_bins=encoder_arg.encoder.n_bins,
-                        encode=encoder_arg.encoder.encode_method,
-                        strategy=encoder_arg.encoder.strategy,
-                        dtype=np.float32,
-                    )
-                    encoded = encoder.fit_transform(
-                        to_encode.to_numpy().reshape(-1, 1)
-                    ).flatten()
-
-                case op_graph.encoder.Binarizer():
-                    from sklearn.preprocessing import Binarizer
-
-                    encoder = Binarizer(
-                        threshold=encoder_arg.encoder.threshold,
-                    )
-                    encoded = encoder.fit_transform(
-                        to_encode.to_numpy().reshape(-1, 1)
-                    ).flatten()
-
-                case op_graph.encoder.MaxAbsScaler():
-                    from sklearn.preprocessing import MaxAbsScaler
-
-                    encoder = MaxAbsScaler()
-                    try:
-                        encoded = encoder.fit_transform(
-                            np.nan_to_num(to_encode.to_numpy()).reshape(-1, 1)
-                        ).flatten()
-                    except ValueError:
-                        encoded = np.array([])
-
-                case op_graph.encoder.StandardScaler():
-                    from sklearn.preprocessing import StandardScaler
-
-                    encoder = StandardScaler(
-                        with_mean=encoder_arg.encoder.with_mean,
-                        with_std=encoder_arg.encoder.with_std,
-                    )
-                    encoded = encoder.fit_transform(
-                        to_encode.to_numpy().reshape(-1, 1)
-                    ).flatten()
-
-                case op_graph.encoder.TimestampEncoder():
-                    if to_encode.dtype == pl.datatypes.Duration:
-                        encoded = self.encode_duration(to_encode)
-                    else:
-                        encoded = self.encode_datetime(to_encode)
-                    source_df = source_df.with_columns(
-                        encoded.rename(encoder_arg.encoded_column_name).cast(
-                            encoder_arg.encoder.output_dtype
-                        )
-                    )
-                    continue
-
-                case op_graph.encoder.TextEncoder():
-                    encoded = self.encode_text(to_encode)
-                    source_df = source_df.with_columns(
-                        encoded.rename(encoder_arg.encoded_column_name).cast(
-                            encoder_arg.encoder.output_dtype
-                        )
-                    )
-                    continue
+            encoded, one_hot_columns = self._encode_column(
+                source_df[encoder_arg.column_name], encoder_arg.encoder
+            )
+            if one_hot_columns is not None:
+                metric[encoder_arg.column_name] = one_hot_columns
 
             source_df = source_df.with_columns(
                 pl.Series(
@@ -1081,7 +1027,7 @@ class InMemoryExecutor(OpGraphExecutor):
             )
         )
 
-    def _execute_embed_node2vec_from_edge_lists(
+    async def _execute_embed_node2vec_from_edge_lists(
         self,
         op: op_graph.op.EmbedNode2vecFromEdgeLists,
         context: _InMemoryExecutionContext,
@@ -1115,7 +1061,7 @@ class InMemoryExecutor(OpGraphExecutor):
 
         edge_list_lfms = list[_LazyFrameWithMetrics]()
         for edge_list in op.edge_list_tables:
-            match self._execute(edge_list.table, context):
+            match await self._execute(edge_list.table, context):
                 case Ok(source_lfm):
                     edge_list_lfms.append(source_lfm)
                 case err:
@@ -1129,58 +1075,65 @@ class InMemoryExecutor(OpGraphExecutor):
                 end_column_type_name = entities_dtypes[end_column_name]
                 metrics.update(lfm.metrics)
                 yield (
-                    lfm.data.with_columns(
-                        pl.col(edge_list.start_column_name).alias(
-                            f"start_id_{start_column_type_name}"
-                        ),
-                        pl.lit(edge_list.start_entity_name).alias("start_source"),
-                        pl.col(edge_list.end_column_name).alias(
-                            f"end_id_{end_column_type_name}"
-                        ),
-                        pl.lit(edge_list.end_entity_name).alias("end_source"),
-                    )
-                    .select(
-                        f"start_id_{start_column_type_name}",
-                        "start_source",
-                        f"end_id_{end_column_type_name}",
-                        "end_source",
+                    context.run_on_worker(
+                        lfm.data.with_columns(
+                            pl.col(edge_list.start_column_name).alias(
+                                f"start_id_{start_column_type_name}"
+                            ),
+                            pl.lit(edge_list.start_entity_name).alias("start_source"),
+                            pl.col(edge_list.end_column_name).alias(
+                                f"end_id_{end_column_type_name}"
+                            ),
+                            pl.lit(edge_list.end_entity_name).alias("end_source"),
+                        )
+                        .select(
+                            f"start_id_{start_column_type_name}",
+                            "start_source",
+                            f"end_id_{end_column_type_name}",
+                            "end_source",
+                        )
+                        .collect
                     )
-                    .collect()
                 )
 
-        edges = pl.concat(
-            [
-                empty_edges_table,
-                *(edge_list for edge_list in edge_generator()),
-            ],
-            rechunk=False,
-            how="diagonal",
-        )
+        async with asyncio.TaskGroup() as tg:
+            edge_tasks = [tg.create_task(edge_list) for edge_list in edge_generator()]
 
-        n2v_space = embed.Space(
-            edges=edges,
-            start_id_column_names=start_id_column_names,
-            end_id_column_names=end_id_column_names,
-            directed=True,
-        )
-        n2v_runner = embed.Node2Vec(
-            space=n2v_space,
-            dim=op.ndim,
-            walk_length=op.walk_length,
-            window=op.window,
-            p=op.p,
-            q=op.q,
-            alpha=op.alpha,
-            min_alpha=op.min_alpha,
-            negative=op.negative,
-        )
-        n2v_runner.train(epochs=op.epochs)
-        return Ok(_LazyFrameWithMetrics(n2v_runner.wv.to_polars().lazy(), metrics))
+        def run_n2v():
+            edges = pl.concat(
+                [
+                    empty_edges_table,
+                    *(task.result() for task in edge_tasks),
+                ],
+                rechunk=False,
+                how="diagonal",
+            )
+            n2v_space = embed.Space(
+                edges=edges,
+                start_id_column_names=start_id_column_names,
+                end_id_column_names=end_id_column_names,
+                directed=True,
+            )
+            n2v_runner = embed.Node2Vec(
+                space=n2v_space,
+                dim=op.ndim,
+                walk_length=op.walk_length,
+                window=op.window,
+                p=op.p,
+                q=op.q,
+                alpha=op.alpha,
+                min_alpha=op.min_alpha,
+                negative=op.negative,
+            )
+            n2v_runner.train(epochs=op.epochs)
+            return n2v_runner.wv.to_polars().lazy()
+
+        return Ok(_LazyFrameWithMetrics(await context.run_on_worker(run_n2v), metrics))
 
-    def _execute_aggregate_columns(
+    async def _execute_aggregate_columns(
         self, op: op_graph.op.AggregateColumns, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
@@ -1205,38 +1158,48 @@ class InMemoryExecutor(OpGraphExecutor):
 
         return Ok(source_lfm.with_data(aggregate))
 
-    def _execute_correlate_columns(
+    async def _execute_correlate_columns(
         self, op: op_graph.op.CorrelateColumns, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
-        with np.errstate(invalid="ignore"):
-            corr_df = source_df.select(op.column_names).corr(dtype="float32")
 
+        def correlate(df: pl.DataFrame):
+            with np.errstate(invalid="ignore"):
+                return df.select(op.column_names).corr(dtype="float32")
+
+        _, corr_df = await context.run_on_worker(
+            _collect_and_apply, source_lfm.data, correlate
+        )
         return Ok(source_lfm.with_data(corr_df.lazy()))
 
-    def _execute_histogram_column(
+    async def _execute_histogram_column(
         self, op: op_graph.op.HistogramColumn, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
-            op.source,
-            lambda lf: lf.collect()[op.column_name]
+        match await self._execute(op.source, context):
+            case Ok(source_lfm):
+                pass
+            case err:
+                return err
+
+        _, result_df = await context.run_on_worker(
+            _collect_and_apply,
+            source_lfm.data,
+            lambda df: df[op.column_name]
             .hist(include_category=False)
-            .lazy()
             .rename(
                 {
                     "breakpoint": op.breakpoint_column_name,
                     "count": op.count_column_name,
                 }
             ),
-            context,
         )
+        return Ok(source_lfm.with_data(result_df.lazy()))
 
-    def _execute_convert_column_to_string(
+    async def _execute_convert_column_to_string(
         self, op: op_graph.op.ConvertColumnToString, context: _InMemoryExecutionContext
     ):
         dtype = op.source.schema.to_polars()[op.column_name]
@@ -1248,14 +1211,15 @@ class InMemoryExecutor(OpGraphExecutor):
             raise NotImplementedError(
                 "converting struct columns to strings is not implemented"
             )
-        return self._compute_source_then_apply(
-            op.source, lambda lf: lf.collect().with_columns(cast_expr).lazy(), context
+
+        return await self._compute_source_then_apply(
+            op.source, lambda lf: lf.with_columns(cast_expr), context
         )
 
-    def _execute_add_row_index(
+    async def _execute_add_row_index(
         self, op: op_graph.op.AddRowIndex, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source,
             lambda lf: lf.with_row_index(
                 name=op.row_index_column_name, offset=op.offset
@@ -1263,70 +1227,76 @@ class InMemoryExecutor(OpGraphExecutor):
             context,
         )
 
-    def _execute_output_csv(
+    async def _execute_output_csv(
         self, op: op_graph.op.OutputCsv, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
-        source_df.write_csv(
-            op.csv_url,
-            quote_style="never",
-            include_header=op.include_header,
+        source_df, _ = await context.run_on_worker(
+            _collect_and_apply,
+            source_lfm.data,
+            lambda df: df.write_csv(
+                op.csv_url, quote_style="never", include_header=op.include_header
+            ),
         )
         return Ok(source_lfm.with_data(source_df.lazy()))
 
-    def _execute_truncate_list(
+    async def _execute_truncate_list(
         self, op: op_graph.op.TruncateList, context: _InMemoryExecutionContext
     ):
         # TODO(Patrick): verify this approach works for arrays
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
-        if len(source_df):
-            existing_length = get_polars_embedding_length(
-                source_df, op.column_name
-            ).unwrap_or_raise()
-        else:
-            existing_length = 0
-        head_length = (
-            op.target_column_length
-            if existing_length >= op.target_column_length
-            else existing_length
-        )
-        source_df = source_df.with_columns(
-            pl.col(op.column_name).list.head(head_length)
+        source_df, existing_length = await context.run_on_worker(
+            _collect_and_apply,
+            source_lfm.data,
+            lambda df: get_polars_embedding_length(df[op.column_name]).unwrap_or_raise()
+            if len(df)
+            else 0,
         )
+
         outer_type = source_df.schema[op.column_name]
         if isinstance(outer_type, pl.Array | pl.List):
             inner_type = outer_type.inner
         else:
             return InternalError("unexpected type", cause="expected list or array type")
+        result = source_df.lazy()
+
+        head_length = (
+            op.target_column_length
+            if existing_length >= op.target_column_length
+            else existing_length
+        )
+        result = result.with_columns(pl.col(op.column_name).list.head(head_length))
 
-        source_df = source_df.lazy()
         if head_length < op.target_column_length:
             padding_length = op.target_column_length - head_length
             padding = [op.padding_value_as_py] * padding_length
-            source_df = source_df.with_columns(
-                pl.col(op.column_name).list.concat(padding)
-            )
-        source_df = source_df.with_columns(
+            result = result.with_columns(pl.col(op.column_name).list.concat(padding))
+        result = result.with_columns(
             pl.col(op.column_name)
             .list.to_array(width=op.target_column_length)
             .cast(pl.List(inner_type))
         )
-        return Ok(source_lfm.with_data(source_df))
+        return Ok(source_lfm.with_data(result))
 
-    def _execute_union(self, op: op_graph.op.Union, context: _InMemoryExecutionContext):
+    async def _execute_union(
+        self, op: op_graph.op.Union, context: _InMemoryExecutionContext
+    ):
+        async with asyncio.TaskGroup() as tg:
+            source_taks = [
+                tg.create_task(self._execute(source, context))
+                for source in op.sources()
+            ]
         sources = list[_LazyFrameWithMetrics]()
-        for source in op.sources():
-            match self._execute(source, context):
+        for task in source_taks:
+            match task.result():
                 case Ok(source_lfm):
                     sources.append(source_lfm)
                 case err:
@@ -1341,16 +1311,19 @@ class InMemoryExecutor(OpGraphExecutor):
             result_lf = result_lf.unique()
         return Ok(_LazyFrameWithMetrics(result_lf, metrics=metrics))
 
-    def _execute_embed_image_column(
+    async def _execute_embed_image_column(
         self, op: op_graph.op.EmbedImageColumn, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
-        to_embed = source_df[op.column_name].cast(pl.Binary())
+        source_df, to_embed = await context.run_on_worker(
+            _collect_and_apply,
+            source_lfm.data,
+            lambda df: df[op.column_name].cast(pl.Binary()),
+        )
 
         embed_context = EmbedImageContext(
             inputs=to_embed,
@@ -1358,43 +1331,38 @@ class InMemoryExecutor(OpGraphExecutor):
             expected_vector_length=op.expected_vector_length,
             expected_coordinate_bitwidth=op.expected_coordinate_bitwidth,
         )
-        match self._image_embedder.embed(embed_context):
+        match await self._image_embedder.aembed(embed_context, context.worker_threads):
             case Ok(result):
                 pass
             case InvalidArgumentError() | InternalError() as err:
                 raise InternalError("Failed to embed column") from err
 
         return Ok(
-            _LazyFrameWithMetrics(
+            source_lfm.with_data(
                 source_df.lazy()
                 .with_columns(result.embeddings.alias(op.embedding_column_name))
-                .drop_nulls(op.embedding_column_name),
-                source_lfm.metrics,
+                .drop_nulls(op.embedding_column_name)
             )
         )
 
-    def _execute_add_decision_tree_summary(
-        self, op: op_graph.op.AddDecisionTreeSummary, context: _InMemoryExecutionContext
+    def _compute_decision_tree_summary(
+        self,
+        data: pl.DataFrame,
+        feature_column_names: Sequence[str],
+        label_column_name: str,
+        max_depth: int,
+        class_names: Sequence[str] | None,
     ):
-        match self._execute(op.source, context):
-            case Ok(source_lfm):
-                pass
-            case err:
-                return err
-
-        df_input = source_lfm.data.collect()
-        dataframe = df_input.select(
-            list({*op.feature_column_names, op.label_column_name})
-        )
+        dataframe = data.select(list({*feature_column_names, label_column_name}))
         boolean_columns = [
             name
             for name, dtype in dataframe.schema.items()
-            if dtype == pl.Boolean() and name in op.feature_column_names
+            if dtype == pl.Boolean() and name in feature_column_names
         ]
 
         # Drop Nan and Null and infinite rows as not supported by decision tree
         dataframe = dataframe.with_columns(
-            *[pl.col(col).cast(pl.Float32) for col in op.feature_column_names]
+            *[pl.col(col).cast(pl.Float32) for col in feature_column_names]
         )
         dataframe = dataframe.drop_nans().drop_nulls()
         try:
@@ -1407,9 +1375,8 @@ class InMemoryExecutor(OpGraphExecutor):
             return InvalidArgumentError(
                 "a minimum of 1 sample is required by DecisionTreeClassifier"
             )
-        features = dataframe[op.feature_column_names]
-        classes = dataframe[op.label_column_name]
-        max_depth = op.max_depth
+        features = dataframe[feature_column_names]
+        classes = dataframe[label_column_name]
 
         from sklearn.tree import DecisionTreeClassifier, export_graphviz, export_text
         from sklearn.utils.multiclass import check_classification_targets
@@ -1427,15 +1394,15 @@ class InMemoryExecutor(OpGraphExecutor):
 
         tree_str = export_text(
             decision_tree=decision_tree,
-            feature_names=op.feature_column_names,
-            class_names=op.classes_names,
+            feature_names=feature_column_names,
+            class_names=class_names,
             max_depth=max_depth,
         )
 
         tree_graphviz = export_graphviz(
             decision_tree=decision_tree,
-            feature_names=op.feature_column_names,
-            class_names=op.classes_names,
+            feature_names=feature_column_names,
+            class_names=class_names,
             max_depth=max_depth,
         )
 
@@ -1445,16 +1412,41 @@ class InMemoryExecutor(OpGraphExecutor):
             )
             tree_str = tree_str.replace(f"{boolean_column} >  0.50", boolean_column)
 
+        return Ok(table_pb2.DecisionTreeSummary(text=tree_str, graphviz=tree_graphviz))
+
+    async def _execute_add_decision_tree_summary(
+        self, op: op_graph.op.AddDecisionTreeSummary, context: _InMemoryExecutionContext
+    ):
+        match await self._execute(op.source, context):
+            case Ok(source_lfm):
+                pass
+            case err:
+                return err
         metrics = source_lfm.metrics.copy()
-        metrics[op.output_metric_key] = table_pb2.DecisionTreeSummary(
-            text=tree_str, graphviz=tree_graphviz
+        source_df, summary_result = await context.run_on_worker(
+            _collect_and_apply,
+            source_lfm.data,
+            lambda df: self._compute_decision_tree_summary(
+                df,
+                op.feature_column_names,
+                op.label_column_name,
+                op.max_depth,
+                op.classes_names,
+            ),
         )
-        return Ok(_LazyFrameWithMetrics(df_input.lazy(), metrics=metrics))
 
-    def _execute_unnest_list(
+        match summary_result:
+            case InvalidArgumentError() | InternalError() as err:
+                return err
+            case Ok(tree_summary):
+                metrics[op.output_metric_key] = tree_summary
+
+        return Ok(_LazyFrameWithMetrics(source_df.lazy(), metrics=metrics))
+
+    async def _execute_unnest_list(
         self, op: op_graph.op.UnnestList, context: _InMemoryExecutionContext
     ):
-        return self._compute_source_then_apply(
+        return await self._compute_source_then_apply(
             op.source,
             lambda lf: lf.with_columns(
                 pl.col(op.list_column_name).list.get(i).alias(column_name)
@@ -1463,47 +1455,42 @@ class InMemoryExecutor(OpGraphExecutor):
             context,
         )
 
-    def _execute_sample_rows(
+    async def _execute_sample_rows(
         self, op: op_graph.op.SampleRows, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
-        n = min(op.num_rows, source_df.shape[0])
-        sample_strategy = op.sample_strategy
-        match sample_strategy:
-            case op_graph.sample_strategy.UniformRandom():
-                result_df = source_df.sample(
-                    n=n,
-                    seed=sample_strategy.seed,
-                )
 
-        return Ok(
-            _LazyFrameWithMetrics(
-                result_df.lazy(),
-                source_lfm.metrics,
-            )
+        def sample(df: pl.DataFrame):
+            match op.sample_strategy:
+                case op_graph.sample_strategy.UniformRandom():
+                    return df.sample(
+                        min(op.num_rows, df.shape[0]), seed=op.sample_strategy.seed
+                    )
+
+        _, result_df = await context.run_on_worker(
+            _collect_and_apply, source_lfm.data, sample
         )
 
-    def _execute_describe_columns(
+        return Ok(_LazyFrameWithMetrics(result_df.lazy(), source_lfm.metrics))
+
+    async def _execute_describe_columns(
         self, op: op_graph.op.DescribeColumns, context: _InMemoryExecutionContext
     ):
-        match self._execute(op.source, context):
+        match await self._execute(op.source, context):
             case Ok(source_lfm):
                 pass
             case err:
                 return err
-        source_df = source_lfm.data.collect()
-        return Ok(
-            source_lfm.with_data(
-                source_df.describe()
-                .lazy()
-                .rename({"statistic": op.statistic_column_name})
-            )
+        _, result_df = await context.run_on_worker(
+            _collect_and_apply,
+            source_lfm.data,
+            lambda df: df.describe().rename({"statistic": op.statistic_column_name}),
         )
+        return Ok(source_lfm.with_data(result_df.lazy()))
 
     def _has_partially_computed_data(
         self, op: op_graph.Op, context: _InMemoryExecutionContext
@@ -1517,7 +1504,7 @@ class InMemoryExecutor(OpGraphExecutor):
             for sub_source in flatten(source.sources() for source in op.sources())
         )
 
-    def _do_execute(  # noqa: C901
+    async def _do_execute(  # noqa: C901
         self,
         op: op_graph.Op,
         context: _InMemoryExecutionContext,
@@ -1552,8 +1539,13 @@ class InMemoryExecutor(OpGraphExecutor):
                     )
                 case Ok(query):
                     pass
-            return self._staging_db.run_select_query(
-                query, expected_schema, context.current_slice_args
+            return (
+                await context.run_on_worker(
+                    self._staging_db.run_select_query,
+                    query,
+                    expected_schema,
+                    context.current_slice_args,
+                )
             ).map(
                 lambda rbr: _LazyFrameWithMetrics(
                     _as_df(rbr, expected_schema).lazy(),
@@ -1571,80 +1563,80 @@ class InMemoryExecutor(OpGraphExecutor):
             case op_graph.op.ReadFromParquet():
                 return self._execute_read_from_parquet(op, context)
             case op_graph.op.RenameColumns():
-                return self._execute_rename_columns(op, context)
+                return await self._execute_rename_columns(op, context)
             case op_graph.op.Join():
-                return self._execute_join(op, context)
+                return await self._execute_join(op, context)
             case op_graph.op.SelectColumns():
-                return self._execute_select_columns(op, context)
+                return await self._execute_select_columns(op, context)
             case op_graph.op.LimitRows():
-                return self._execute_limit_rows(op, context)
+                return await self._execute_limit_rows(op, context)
             case op_graph.op.OffsetRows():
-                return self._execute_offset_rows(op, context)
+                return await self._execute_offset_rows(op, context)
             case op_graph.op.OrderBy():
-                return self._execute_order_by(op, context)
+                return await self._execute_order_by(op, context)
             case op_graph.op.FilterRows():
-                return self._execute_filter_rows(op, context)
+                return await self._execute_filter_rows(op, context)
             case op_graph.op.DistinctRows():
-                return self._execute_distinct_rows(op, context)
+                return await self._execute_distinct_rows(op, context)
             case (
                 op_graph.op.SetMetadata()
                 | op_graph.op.UpdateMetadata()
                 | op_graph.op.RemoveFromMetadata()
                 | op_graph.op.UpdateFeatureTypes()
             ):
-                return self._execute(op.source, context)
+                return await self._execute(op.source, context)
             case op_graph.op.EmbeddingMetrics() as op:
-                return self._execute_embedding_metrics(op, context)
+                return await self._execute_embedding_metrics(op, context)
             case op_graph.op.EmbeddingCoordinates():
-                return self._execute_embedding_coordinates(op, context)
+                return await self._execute_embedding_coordinates(op, context)
             case op_graph.op.RollupByAggregation() as op:
-                return self._execute_rollup_by_aggregation(op, context)
+                return await self._execute_rollup_by_aggregation(op, context)
             case op_graph.op.Empty():
-                return self._execute_empty(op, context)
+                return await self._execute_empty(op, context)
             case op_graph.op.EmbedNode2vecFromEdgeLists():
-                return self._execute_embed_node2vec_from_edge_lists(op, context)
+                return await self._execute_embed_node2vec_from_edge_lists(op, context)
             case op_graph.op.Concat():
-                return self._execute_concat(op, context)
+                return await self._execute_concat(op, context)
             case op_graph.op.UnnestStruct():
-                return self._execute_unnest_struct(op, context)
+                return await self._execute_unnest_struct(op, context)
             case op_graph.op.NestIntoStruct():
-                return self._execute_nest_into_struct(op, context)
+                return await self._execute_nest_into_struct(op, context)
             case op_graph.op.AddLiteralColumn():
-                return self._execute_add_literal_column(op, context)
+                return await self._execute_add_literal_column(op, context)
             case op_graph.op.CombineColumns():
-                return self._execute_combine_columns(op, context)
+                return await self._execute_combine_columns(op, context)
             case op_graph.op.EmbedColumn():
-                return self._execute_embed_column(op, context)
+                return await self._execute_embed_column(op, context)
             case op_graph.op.EncodeColumns():
-                return self._execute_encode_columns(op, context)
+                return await self._execute_encode_columns(op, context)
             case op_graph.op.AggregateColumns():
-                return self._execute_aggregate_columns(op, context)
+                return await self._execute_aggregate_columns(op, context)
             case op_graph.op.CorrelateColumns():
-                return self._execute_correlate_columns(op, context)
+                return await self._execute_correlate_columns(op, context)
             case op_graph.op.HistogramColumn():
-                return self._execute_histogram_column(op, context)
+                return await self._execute_histogram_column(op, context)
             case op_graph.op.ConvertColumnToString():
-                return self._execute_convert_column_to_string(op, context)
+                return await self._execute_convert_column_to_string(op, context)
             case op_graph.op.AddRowIndex():
-                return self._execute_add_row_index(op, context)
+                return await self._execute_add_row_index(op, context)
             case op_graph.op.OutputCsv():
-                return self._execute_output_csv(op, context)
+                return await self._execute_output_csv(op, context)
             case op_graph.op.TruncateList():
-                return self._execute_truncate_list(op, context)
+                return await self._execute_truncate_list(op, context)
             case op_graph.op.Union():
-                return self._execute_union(op, context)
+                return await self._execute_union(op, context)
             case op_graph.op.EmbedImageColumn():
-                return self._execute_embed_image_column(op, context)
+                return await self._execute_embed_image_column(op, context)
             case op_graph.op.AddDecisionTreeSummary():
-                return self._execute_add_decision_tree_summary(op, context)
+                return await self._execute_add_decision_tree_summary(op, context)
             case op_graph.op.UnnestList():
-                return self._execute_unnest_list(op, context)
+                return await self._execute_unnest_list(op, context)
             case op_graph.op.SampleRows():
-                return self._execute_sample_rows(op, context)
+                return await self._execute_sample_rows(op, context)
             case op_graph.op.DescribeColumns():
-                return self._execute_describe_columns(op, context)
+                return await self._execute_describe_columns(op, context)
 
-    def _execute(
+    async def _execute(
         self,
         op: op_graph.Op,
         context: _InMemoryExecutionContext,
@@ -1681,7 +1673,7 @@ class InMemoryExecutor(OpGraphExecutor):
 
             try:
                 _logger.info("starting op execution")
-                maybe_lfm = self._do_execute(op=op, context=context)
+                maybe_lfm = await self._do_execute(op=op, context=context)
             finally:
                 _logger.info("op execution complete")
             match maybe_lfm:
@@ -1703,29 +1695,33 @@ class InMemoryExecutor(OpGraphExecutor):
                 context.computed_batches_for_op_graph[sliced_table] = lfm
             return Ok(lfm)
 
-    def execute(
-        self, context: ExecutionContext
+    async def execute(
+        self,
+        context: ExecutionContext,
+        worker_threads: ThreadPoolExecutor | None = None,
     ) -> (
         Ok[ExecutionResult]
         | InvalidArgumentError
         | InternalError
         | ResourceExhaustedError
     ):
-        with _InMemoryExecutionContext(context) as in_memory_context:
+        with _InMemoryExecutionContext(context, worker_threads) as in_memory_context:
             for table_context in context.tables_to_compute:
                 in_memory_context.current_output_context = table_context
                 sliced_table = _SlicedTable(
                     table_context.table_op_graph, table_context.sql_output_slice_args
                 )
                 if sliced_table not in in_memory_context.computed_batches_for_op_graph:
-                    match self._execute(sliced_table.op_graph, in_memory_context):
+                    match await self._execute(sliced_table.op_graph, in_memory_context):
                         case Ok():
                             pass
                         case err:
                             return err
         args_lfm_iterator = in_memory_context.computed_batches_for_op_graph.items()
         computed_tables = {
-            slice_args: _SchemaAndBatches.from_lazy_frame_with_metrics(lfm)
+            slice_args: await _SchemaAndBatches.from_lazy_frame_with_metrics(
+                lfm, worker_threads
+            )
             for slice_args, lfm in args_lfm_iterator
         }