composability.py

import xorq.api as xo
from xorq.vendor import ibis
from xorq.expr.udf import agg
import xorq.expr.datatypes as dt

diamonds = xo.examples.diamonds.fetch().tag("source")


def price_stats(df):
    """
    Shows basic statistical aggregation pattern.
    """
    import numpy as np

    prices = df['price'].dropna()

    if len(prices) == 0:
        return {'median': 0.0, 'std': 0.0, 'iqr': 0.0}

    return {
        'median': float(np.median(prices)),
        'std': float(np.std(prices)),
        'iqr': float(np.percentile(prices, 75) - np.percentile(prices, 25))
    }


price_stats_schema = diamonds.select(['price']).schema()

price_stats_udaf = agg.pandas_df(
    fn=price_stats,
    schema=price_stats_schema,
    return_type=dt.Struct({
        'median': dt.float64,
        'std': dt.float64,
        'iqr': dt.float64
    }),
    name="price_stats"
)


high_value_quality = (
    diamonds
    .filter(diamonds.price > 1000)  # Filter first
    .group_by('cut')
    .agg(
        stats=price_stats_udaf.on_expr(diamonds)
    )
)


expr = high_value_quality

"""
❯ xorq build scripts/composability.py -e diamonds
Building diamonds from scripts/composability.py
/Users/hussainsultan/workspace/xorq/python/xorq/ibis_yaml/translate.py:560: UserWarning: The Read op path is using a local filesystem path, running the build may not work in other environments.
  warnings.warn(
Written 'diamonds' to builds/17d96efb38c0
builds/17d96efb38c0

xorq main*​​ ≡3s
❯ xorq build scripts/composability.py -e expr
Building expr from scripts/composability.py
/Users/hussainsultan/workspace/xorq/python/xorq/ibis_yaml/translate.py:560: UserWarning: The Read op path is using a local filesystem path, running the build may not work in other environments.
  warnings.warn(
Written 'expr' to builds/4a90ac01bb90
builds/4a90ac01bb90

❯  xorq run builds/17d96efb38c0 -f arrow -o - | xorq run-unbound builds/4a90ac01bb90 --to_unbind_tag "source" -f arrow -o - | duckdb -c "load arrow;
  select cut, unnest(stats) from read_arrow('/dev/stdin');"

[run-unbound] Loading expression from builds/4a90ac01bb90
[run-unbound] Reading Arrow IPC from instream...
┌───────────┬────────┬───────────────────┬─────────┐
│    cut    │ median │        std        │   iqr   │
│  varchar  │ double │      double       │ double  │
├───────────┼────────┼───────────────────┼─────────┤
│ Good      │ 3965.0 │ 3664.345196837687 │  3676.5 │
│ Fair      │ 3564.0 │ 3547.398745542728 │  3179.5 │
│ Very Good │ 4136.0 │ 3949.144001457176 │ 4372.25 │
│ Premium   │ 4562.0 │ 4334.525176388644 │  5326.0 │
│ Ideal     │ 3161.0 │ 4022.340182823431 │  4678.0 │
└───────────┴────────┴───────────────────┴─────────┘
"""