Read Primitives — Optimizer: Union Optimization

Read Primitives benchmark · all Read Primitives queries

Test union optimization with multiple scans

Representative SQL

Rendered for DataFusion with default parameters.

SELECT
  c_name,
  c_mktsegment,
  'high_value' AS customer_type,
  c_acctbal
FROM customer
WHERE
  c_acctbal > 8000 AND c_nationkey IN (1, 2, 3)
UNION ALL
SELECT
  c_name,
  c_mktsegment,
  'medium_value' AS customer_type,
  c_acctbal
FROM customer
WHERE
  c_acctbal BETWEEN 4000 AND 8000 AND c_nationkey IN (1, 2, 3)
UNION ALL
SELECT
  c_name,
  c_mktsegment,
  'low_value' AS customer_type,
  c_acctbal
FROM customer
WHERE
  c_acctbal BETWEEN 1000 AND 4000 AND c_nationkey IN (1, 2, 3)
ORDER BY
  c_acctbal DESC
LIMIT 300

Representative DataFrame

Expression family (Polars, PySpark, DataFusion), default parameters.

def optimizer_union_optimization_expression_impl(ctx: DataFrameContext) -> Any:
    col = ctx.col
    lit = ctx.lit

    customer = ctx.get_table("customer")
    in_nations = col("c_nationkey").is_in([1, 2, 3])
    out_cols = ("c_name", "c_mktsegment", "customer_type", "c_acctbal")
    high = (
        customer.filter((col("c_acctbal") > lit(8000)) & in_nations)
        .with_columns(lit("high_value").alias("customer_type"))
        .select(*out_cols)
    )
    medium = (
        customer.filter((col("c_acctbal") >= lit(4000)) & (col("c_acctbal") <= lit(8000)) & in_nations)
        .with_columns(lit("medium_value").alias("customer_type"))
        .select(*out_cols)
    )
    low = (
        customer.filter((col("c_acctbal") >= lit(1000)) & (col("c_acctbal") <= lit(4000)) & in_nations)
        .with_columns(lit("low_value").alias("customer_type"))
        .select(*out_cols)
    )
    return ctx.concat([high, medium, low]).sort(["c_acctbal", "c_name"], descending=[True, False]).limit(300)

Get the exact query for your run

The blocks above are representative. To get the precise statement a specific platform runs at a specific scale, substitute your own values:

from benchbox import ReadPrimitives

bench = ReadPrimitives(scale_factor=0.01)            # your scale factor
queries = bench.get_queries(dialect="datafusion")   # your platform's dialect; drop dialect= for native SQL
print(queries['optimizer_union_optimization'])

For the DataFrame form:

from benchbox.core.query_catalog import get_dataframe_render

print(get_dataframe_render('read_primitives', 'optimizer_union_optimization', family="expression").source)  # or family="pandas"

Source

Query text: benchbox/core/read_primitives/queries.py