raiju 0.1.3

Distributed PySpark execution layer for composable data engineering workflows (Spark-native, extension-ready)

Raiju

Docs

Raiju is a distributed PySpark execution framework aimed at teams who need maintainable orchestration for data engineering workflows that stretch beyond simple columnar transforms—while staying Spark-native on the cluster.

In one line: Raiju is built on PySpark to simplify complex transformation orchestration, scalable UDF-style workflows, and (as the project grows) integrated local and remote inference for operational enrichment—without replacing Spark.

Raiju (雷獣, raijū) is a creature from Japanese folklore: a lightning beast and companion of the thunder god Raijin. The name fits a layer that rides PySpark, your engine for distributed data processing.

Table of contents

1. Why Raiju exists
2. Core design goals
3. Architecture overview
4. Integrated inference workflows (direction)
5. Example workloads
6. What you get today
7. Roadmap
8. Getting started
9. Inference settings (Ollama / OpenRouter)
10. Weave (broadcast-friendly joins)
11. Weft (LLM schema aliasing and typing)
12. DataFrame profiling
13. How it works
14. Development
15. Changelog
16. Support
17. Show your support
18. Code of conduct
19. Contributing
20. License
21. Security

Why Raiju exists

Many distributed data workflows are hard to keep healthy when:

• transformations need cross-record or multi-field reasoning
• logic becomes deeply procedural across many steps
• orchestration spans several enrichment or validation stages
• you need dynamic execution choices without scattering if trees through jobs
• column-only pipelines are hard to read or refactor
• UDF-heavy pipelines become brittle operationally

PySpark already gives you distributed compute. What teams often lack is a clear, composable layer for orchestration and advanced workflows—without giving up executors, partitions, and the rest of the Spark programming model. Raiju is meant to grow into that layer: higher-level workflow composition on top of unchanged Spark execution.

Core design goals

• Distributed-first: Spark executors and cluster semantics stay central; Raiju coordinates and composes, it does not pretend compute is “local-first.”
• Workflow composition: Modular pipelines and reusable building blocks instead of one-off scripts.
• Operational flexibility: Room for local inference, remote providers, and hybrid patterns where privacy, cost, or latency demand it.
• Developer ergonomics: Less bespoke glue for complex jobs; clearer boundaries between stages.
• Spark compatibility: Same DataFrame types, same session APIs, same deployment story (including Databricks and on-prem clusters).

Architecture overview

Target shape: Raiju sits as an orchestration and execution abstraction above Spark-native processing—handling transformation composition, coordination patterns, enrichment flows, and (optionally) inference calls—while work still runs on Spark executors.

Today: the published library is a thin, delegation-based SparkSession entry point (see What you get today). That is intentional: a stable compatibility surface before higher-level APIs land. The layering story above is where the project is headed; see ROADMAP.md for concrete backlog items (diagrams, retries, partitioning docs, inference interfaces, benchmarks).

Integrated inference workflows (direction)

Optional support for LLM-assisted steps inside data pipelines is part of the vision—framed as operational enrichment, not a separate “agent platform.”

Planned execution styles to document and implement over time:

• Local inference (for example via Ollama) for low-latency or air-gapped settings
• Remote providers (for example OpenRouter-compatible HTTP APIs) when external models are acceptable
• Hybrid routing by policy (cost, privacy, SLA)

Example workload types (all squarely “data engineering”):

• semantic enrichment and tagging
• entity normalization and fuzzy classification
• metadata generation and schema hints
• human-in-the-loop review assistance as a batch step
• contextual transforms where a model proposes a value validated by rules

Language in the project intentionally stays grounded: orchestration, enrichment, inference hooks—not hype around autonomy or “cognitive” stacks.

You can initialize Raiju with provider settings (endpoints, default models, OpenRouter key resolution) so later orchestration can call into Ollama or OpenRouter without ad-hoc globals. Constructing InferenceSettings performs no HTTP requests—it only holds configuration. Bounded driver-side HTTP runs when you opt in to DataFrame profiling LLM enrichment or Weft canonicalization—both use Pydantic-validated model JSON, not free-form parsing. The same transport and JSON-object parsing live in raiju.inference.chat and are re-exported from raiju as inference_chat, parse_llm_json_object, and truncate_llm_text for custom tools or tests.

Example workloads

Raiju is aimed at teams building operational data systems where jobs look like:

• large-scale semantic enrichment
• complex mapPartitions / UDF orchestration
• fuzzy entity resolution and deduplication
• metadata standardization across sources
• multi-stage enrichment with checkpoints
• operational anomaly or triage classification
• hybrid rules + model scoring in batch
• dataset QA: column-level profiles with optional LLM regex and format hints before publishing

What you get today

Release v0.1.3 ships a single, extension-ready entry point over PySpark plus profiling utilities:

• Full PySpark surface: Raiju forwards the entire SparkSession API via delegation—no duplicated method lists; new PySpark APIs keep working as PySpark evolves.
• Drop-in usage: Raiju.builder...getOrCreate() or Raiju(spark) when you already have a session (for example in Databricks).
• Inference settings on the session: optional InferenceSettings (Ollama and/or OpenRouter) attached at construction or via with_inference() for builder flows (configuration only at init).
• weave joins: optional broadcast() hints when one side is much smaller than the other, using bounded row-count inference (see Weave (broadcast-friendly joins)).
• Weft (schema prep): weft_dataframe / Raiju.weft map messy source columns onto a canonical dict-of-fields you define, using one bounded LLM call plus Pydantic (WeftResponse / WeftColumnMapping), confidence guardrails, optional Spark-native casts (single select), optional python-dateutil fuzzy timestamps when the model requests that path, and optional struct output. See Weft (LLM schema aliasing and typing).
• DataFrame profiling: profile_dataframe / Raiju.profile compute rich per-column stats in Spark-native aggregates (optional freqItems, optional bounded LLM enrichment with Pydantic-validated JSON and tiktoken token estimates). See DataFrame profiling.
• Inference transport (library API): inference_chat, parse_llm_json_object, and truncate_llm_text from raiju (implemented in raiju.inference.chat) — same Ollama/OpenRouter routing, timeouts, and JSON fence handling Weft and profiling use; useful for bespoke prompts outside those entry points.
• Runtime dependencies: pyspark>=4.0, pydantic>=2.5, tiktoken>=0.7, python-dateutil>=2.8 (dateutil is used when Weft applies a model-requested fuzzy-parse strategy; profiling does not require it for core stats).

Higher-level orchestration beyond Weft, profiling, weave, generic HTTP inference clients, and operational guides are on the roadmap (ROADMAP.md).

Roadmap

See ROADMAP.md for a structured backlog: execution and DAG diagrams, failure handling and retry semantics, partitioning and serialization notes, benchmarks, orchestration APIs, optional inference backends, and hardening for production pipelines.

Getting started

Installation

Raiju is published as raiju on PyPI.

With uv (recommended), pip, or pipx:

# With uv.
uv add raiju                    # Add to your project.
uv tool install raiju@latest    # Or install globally.

# With pip.
pip install raiju

# With pipx.
pipx install raiju

From a local clone:

pip install -e .

For development (linting, formatting, tests):

pip install -e ".[dev]"

Requirements: Python 3.9+, PySpark 4.0+, plus pydantic, tiktoken, and python-dateutil (pulled in automatically with pip install raiju).

Usage

Create a session with the builder:

from raiju import Raiju

raiju = Raiju.builder.appName("my_app").master("local[*]").getOrCreate()

Or wrap an existing session (for example in Databricks):

from raiju import Raiju

raiju = Raiju(spark)

Use it like PySpark: SQL, DataFrame API, read, catalog, config. Everything is delegated:

# SQL
df = raiju.sql("SELECT 1 AS one")

# DataFrame API
df = raiju.range(10).filter("id > 5")

# Read data
df = raiju.read.csv("path/to/file.csv", header=True)

# Catalog, UDFs, config
raiju.catalog.listTables()
raiju.conf.set("key", "value")

Returned objects are standard PySpark types.

Quick profile (Spark-side stats only; see DataFrame profiling for LLM options):

from raiju import ProfileOptions, profile_dataframe

sample = raiju.range(1_000).toDF("id")
print(profile_dataframe(sample, options=ProfileOptions())["row_count"])

Inference settings (Ollama / OpenRouter)

Attach one or both backends so profiling LLM enrichment and future Raiju execution can read models and endpoints from raiju.inference (no network I/O at init):

from pyspark.sql import SparkSession
from raiju import InferenceSettings, OllamaConfig, OpenRouterConfig, Raiju

spark = SparkSession.builder.appName("enrich").master("local[*]").getOrCreate()
raiju = Raiju(
    spark,
    inference=InferenceSettings(
        ollama=OllamaConfig(default_model="llama3.2"),
        openrouter=OpenRouterConfig(
            default_model="anthropic/claude-3.5-sonnet",
            # api_key=None → reads OPENROUTER_API_KEY from the environment (emits a UserWarning once at config construction)
        ),
    ),
)

assert raiju.inference is not None
assert raiju.inference.ollama.default_model == "llama3.2"
assert raiju.inference.openrouter.resolved_api_key() is not None  # if env is set

If you use Raiju.builder...getOrCreate(), the builder still returns a bare Raiju; chain with_inference(...) on the result (same underlying SparkSession):

from raiju import InferenceSettings, OllamaConfig, Raiju

raiju = Raiju.builder.appName("enrich").master("local[*]").getOrCreate().with_inference(
    InferenceSettings(ollama=OllamaConfig(default_model="llama3.2"))
)

Weave (broadcast-friendly joins)

weave joins two PySpark DataFrames and can attach Spark’s broadcast() hint when one relation is clearly much smaller than the other. The decision uses bounded counts—at most bounded_count_cap + 1 rows are counted per side—so you avoid a full count() on huge tables while still getting an explicit hint when the skew is obvious.

Spark’s planner already auto-broadcasts when estimated size is under spark.sql.autoBroadcastJoinThreshold. Use weave when statistics are missing or conservative and you still want a deliberate broadcast join path.

Raiju.weave

On a Raiju session, weave delegates to the same join logic as the module function (bounded counts and broadcast() hints only—no inference or LLM calls):

from pyspark.sql import SparkSession
from raiju import Raiju

spark = SparkSession.builder.appName("demo").master("local[*]").getOrCreate()
raiju = Raiju(spark)

result = raiju.weave(large_df, small_df, on="id", how="left_outer")

With the builder:

from raiju import Raiju

raiju = Raiju.builder.appName("demo").master("local[*]").getOrCreate()
out = raiju.weave(facts, dim, on="sk", how="inner")

weave from the module

The same behavior is available without a Raiju wrapper—useful in shared libraries or plain SparkSession code:

from raiju import BroadcastJoinPolicy, weave

out = weave(
    large_df,
    small_df,
    on="id",
    how="inner",
    policy=BroadcastJoinPolicy(
        bounded_count_cap=50_000,
        max_small_to_large_ratio=0.15,
    ),
)

broadcast_side

Value	Behavior
"auto" (default)	Infer which side to broadcast from bounded counts and BroadcastJoinPolicy.
"left" / "right"	Force broadcast() on that operand.
"none"	Ordinary join with no broadcast hint.

BroadcastJoinPolicy fields

Field	Default	Role
bounded_count_cap	100_000	Each side uses limit(cap + 1).count() so at most cap + 1 rows are read for the decision.
max_small_to_large_ratio	0.2	Broadcast only if the smaller bounded count is at most this fraction of the larger.
ambiguous_when_both_at_cap	True	If both sides hit cap + 1, skip broadcasting (sizes are unclear).

Tune these together with executor memory and spark.sql.autoBroadcastJoinThreshold so broadcast joins stay within cluster limits.

Weft (LLM schema aliasing and typing)

Weft aligns a PySpark DataFrame to a canonical schema you describe: keys are target column names, values are natural-language field descriptions (examples, semantics, edge cases). A single bounded LLM request (aggregates + capped samples per column, similar spirit to profiling enrichment) proposes rename + typing decisions. Raiju validates the reply with Pydantic, applies confidence guardrails, then builds the result in one Spark select from the original frame—no withColumn chains for the canonical block.

Typical flow: Weft normalizes column semantics and types; Weave can join normalized datasets. Together they are the “prep + link” story for messy sources.

Requirements

• InferenceSettings on the session (or passed explicitly) with Ollama and/or OpenRouter—same configuration model as profiling enrichment.
• A non-empty structure mapping: canonical_name → description string.

Raiju.weft and weft_dataframe

from raiju import InferenceSettings, OllamaConfig, Raiju

raiju = Raiju(spark, inference=InferenceSettings(ollama=OllamaConfig(default_model="llama3.2")))

structure = {
    "payee_name": "Entity receiving payment; vendor or person name.",
    "payment_amount": "Numeric payment amount; may include currency symbols in source.",
    "payment_date": "Date the payment was recorded; many string formats possible.",
}

mapped, report = raiju.weft(
    df,
    structure,
    min_confidence=0.85,
    require_review_below=0.95,
    allow_unmapped=False,
    return_report=True,
)
# report: accepted_mappings, accepted_specs, typing_applied, nullability_applied,
# needs_review, review_suggested, ignored_columns, confidence_scores,
# weft_advisory, struct_schema_simple (if output="struct"), llm_token_usage, …

The module function weft_dataframe(df, structure, inference, ...) is the same API without a Raiju wrapper. For tests or custom pipelines, resolve_weft_mappings applies guardrails to an already-validated WeftResponse without HTTP.

Main options

Parameter	Default	Role
min_confidence	0.85	Below this, a model map is not applied (column stays unresolved unless ignore).
require_review_below	0.95	Applied maps in [min_confidence, require_review_below) are listed under report["review_suggested"].
allow_unmapped	False	If False, error when a source column is neither accepted nor explicitly ignored after rules.
allow_many_to_one	False	If False, multiple sources mapping to one target are withheld and flagged; if True, highest-confidence source wins, others flagged.
apply_typing	True	Coerce to target_spark_type from the model (strings → safe numeric/temporal paths, try_to_timestamp coalesce, optional dateutil UDF when requested).
output	"flat"	"flat" — one column per canonical key in structure order. "struct" — nest them under struct_name.
struct_name	"weft"	Struct column name when output="struct".
keep_extra_columns	False	Append unmapped, non-ignored source columns after the canonical block (still one select).
emit_weft_warnings	True	Emit WeftWarning for missing canonical slots, dateutil throughput, default date formats, etc.
provider	"auto"	"ollama", "openrouter", or "auto" (prefer Ollama if configured); normalized case-insensitively.
http_timeout_s	120	Per-request HTTP timeout for the Weft LLM call (must be positive).
sample_scan_limit / max_sample_values / max_value_chars	120 / 14 / 280	Bound the evidence payload sent to the model.

Pydantic contract (Weft)

Import from raiju or raiju.inference:

• WeftResponse — mappings, unmapped_columns, ambiguous_columns, notes.
• WeftColumnMapping — per source column: target_column, confidence, reason, action (map | ignore | needs_review), target_spark_type, nullable, decimal precision/scale, temporal_parse_strategy (native | spark_formats | python_dateutil), spark_timestamp_formats, python_dateutil_fuzzy.
• WeftWarning — advisories for casts, missing canonical fields, and fuzzy date UDF use (filter with warnings.filterwarnings).

Performance and safety notes

• Single scan for the canonical projection: renames and casts are expressed as one select of column expressions (plus optional extra columns).
• dateutil runs inside a Python UDF only when the validated model output sets temporal_parse_strategy to python_dateutil; prefer spark_formats with explicit patterns when possible.
• Token usage: successful Weft HTTP calls attach llm_token_usage to the report and emit RaijuLLMUsageWarning the same way profiling enrichment does.

DataFrame profiling

profile_dataframe (and Raiju.profile) summarize a PySpark DataFrame with Spark-native aggregates—one wide agg over the input for most metrics, plus optional freqItems and optional driver-side LLM enrichment when you attach InferenceSettings and opt in. This path is built for throughput: it does not scan the table in a Python row UDF to compute column stats.

What you get

• Row and column metrics: null counts, completeness, approximate distinct counts, type-aware stats (numeric percentiles / skew / kurtosis, string length and lexicographic bounds, booleans, timestamps, arrays, maps, and more).
• describe-style blocks under each column where applicable, plus structured JSON-friendly output (collect=True by default scrubs NaN/inf for logging).
• Optional freqItems: approximate frequent values / mode-style candidates (ProfileOptions.include_freq_items).
• Optional LLM layer: when inference is InferenceSettings and inference_enrichment=True, Raiju sends aggregates plus a capped row sample (single bounded collect) to Ollama or OpenRouter. The model returns JSON validated by Pydantic (ProfileEnrichmentResponse). Suggested regexes are checked with Python re.compile; failures become UserWarnings and null llm payloads.
• Token accounting: prompt and completion sizes are estimated with tiktoken (build_llm_token_usage); provider-reported usage is copied under raw_usage["api"] for audit only. After a successful HTTP call, Raiju emits RaijuLLMUsageWarning with the tiktoken totals (filter with warnings.filterwarnings if needed).

profile_dataframe (function API)

from raiju import ProfileOptions, profile_dataframe

profile = profile_dataframe(
    df,
    options=ProfileOptions(
        percentiles=[0.25, 0.5, 0.75],
        include_freq_items=True,
        freq_items_support=0.05,
        columns=["user_id", "event_ts", "payload"],  # optional subset
    ),
)
print(profile["row_count"])
print(profile["columns"]["user_id"]["approx_distinct"])
print(profile["approximate_spark_actions"])  # 1 + 1 if freqItems ran

Raiju.profile (session shortcut)

Same as profile_dataframe, and forwards Raiju.inference when you omit inference=:

from raiju import InferenceSettings, OllamaConfig, ProfileOptions, Raiju

raiju = Raiju(spark, inference=InferenceSettings(ollama=OllamaConfig(default_model="llama3.2")))
prof = raiju.profile(df, options=ProfileOptions(inference_enrichment=True))

LLM enrichment options (ProfileOptions)

Field	Default	Role
inference_enrichment	False	When True and inference is InferenceSettings, run the bounded LLM pass.
inference_provider	"auto"	"auto" prefers Ollama if configured, else OpenRouter; or force "ollama" / "openrouter".
inference_http_timeout_s	120	HTTP timeout for the chat request.
inference_max_columns	28	Max columns sent to the model (prioritizes strings, temporal, boolean, then low-cardinality numerics).
inference_sample_scan_limit	120	Max rows read for the batched sample select.
inference_max_sample_values	14	Max distinct sample strings/values per column in the payload.
inference_max_value_chars	280	Truncate sample cell text for the prompt.

Pydantic types (LLM JSON contract)

Import from raiju or raiju.inference:

• ProfileEnrichmentResponse — root object with columns: list[ProfileEnrichmentColumn].
• ProfileEnrichmentColumn — one column’s LLM fields (human_summary, suggested_validation_regex, java_simple_date_format, python_strptime_directive, pii_likelihood, …).
• WeftResponse / WeftColumnMapping — Weft LLM JSON contract (see Weft).
• WeftWarning — Weft-specific advisory warnings.
• LLMTokenUsage — normalized token fields plus raw_usage (tiktoken metadata + api blob).

profile_to_describe_rows

Flatten describe-compatible entries into row records (for notebooks or small tables):

from raiju import profile_dataframe, profile_to_describe_rows

prof = profile_dataframe(df)
rows = profile_to_describe_rows(prof)
# each row: {"summary": "mean"|"count"|..., "column": str, "value": ...}

OpenRouter example (explicit API key)

import warnings

from raiju import (
    InferenceSettings,
    OpenRouterConfig,
    ProfileOptions,
    RaijuLLMUsageWarning,
    profile_dataframe,
)

warnings.simplefilter("always", RaijuLLMUsageWarning)

inf = InferenceSettings(
    openrouter=OpenRouterConfig(
        api_key="sk-or-...",  # or None + OPENROUTER_API_KEY
        default_model="openai/gpt-4o-mini",
    ),
)

prof = profile_dataframe(
    df,
    inference=inf,
    options=ProfileOptions(
        inference_enrichment=True,
        inference_provider="openrouter",
        inference_max_columns=12,
    ),
)

assert prof.get("llm_enrichment", {}).get("status") in ("ok", "failed")
print(prof.get("llm_token_usage"))  # tiktoken-based totals when HTTP succeeded

Advanced: reuse the same chat transport

Weft and profiling both call inference_chat under the hood. For a custom system prompt and user payload (still one bounded request), use the same API and token accounting:

from raiju import InferenceSettings, OllamaConfig, inference_chat

inf = InferenceSettings(ollama=OllamaConfig(default_model="llama3.2"))
used, text, usage = inference_chat(
    inf,
    system="You reply with a single JSON object only.",
    user='{"task": "ping"}',
    provider="auto",
    http_timeout_s=60.0,
    purpose="my_tool",
)
# `usage` is normally present; call usage.warn_if_known() like Weft does after a successful call.

parse_llm_json_object strips optional Markdown JSON fences and returns a dict or None if the model did not return a JSON object. truncate_llm_text caps strings for prompts and logs.

Advanced: count tokens outside profiling

If you have raw strings (for example from another tool), you can reuse the same accounting Raiju uses after HTTP:

from raiju import build_llm_token_usage

usage = build_llm_token_usage(
    provider="openrouter",
    model="openai/gpt-4o-mini",
    system="You are a helpful assistant.",
    user='{"task": "summarize"}',
    assistant='{"ok": true}',
    raw_api_response={},  # optional provider JSON to stash under raw_usage["api"]
)
print(usage.total_tokens, usage.model_dump()["raw_usage"]["tiktoken"]["encoding"])

Notes

• Cost and privacy: LLM enrichment sends bounded aggregates and samples only, but still leaves your network boundary—treat models and prompts like production data.
• Spark actions: Profiling always runs at least one aggregate action; freqItems and LLM sampling add additional actions documented in approximate_spark_actions.

How it works

• No hardcoded API surface: Raiju and its builder use __getattr__ to forward to the real SparkSession (and SparkSession.builder).
• Single entry point: You hold a Raiju instance; .read, .sql, .range, and the rest behave as in PySpark.
• Composable helpers: weave, weft, profile / profile_dataframe, weft_dataframe / resolve_weft_mappings, and inference_chat / parse_llm_json_object / truncate_llm_text live on Raiju or as module functions without replacing DataFrame types.
• Thin foundation: This layer is the base for future orchestration and enrichment utilities without forking PySpark.

Development

pip install -e ".[dev]"
ruff check raiju/ tests/
ruff format raiju/ tests/
pytest tests/ -v

Changelog

See CHANGELOG.md for release history.

Support

Having trouble? Open an issue on GitHub.

Code of conduct

This project adheres to the Contributor Covenant Code of Conduct. By participating, you are expected to uphold this code.

Contributing

Contributions are welcome. See CONTRIBUTING.md for how to get started.

Show your support

If you are using Raiju, consider adding the Raiju badge to your project’s README.md:

[![Raiju](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/seanpavlak/raiju/main/assets/badge.json)](https://github.com/seanpavlak/raiju)

License

This repository is licensed under the MIT License.

Security

To report a security concern or vulnerability, see SECURITY.md.