hexadruid 0.1.8

Smart Spark Optimizer: Skew Rebalancer + Key Detector + DRTree

HexaDruid 🧠⚡

HexaDruid is an intelligent Spark optimizer designed to tackle data skew, ambiguous key detection, and schema bloat using smart salting, recursive shard-aware rule trees, and adaptive tuning. It enables better parallelism, safer memory layout, and intelligent insight into skewed datasets using PySpark’s native DataFrame API.

---

🚀 Installation

pip install hexadruid

To upgrade to the latest version?

pip install --upgrade hexadruid

---

🔍 Features

• 📊 Smart Salting using Z-score or IQR skew analysis + percentile bucketing
• 🌲 Recursive DRTree for shard-based logical filtering with SQL predicates
• 🔑 Primary & Composite Key Detection (UUIDs, alphanumerics, hex — optional)
• 🧠 Schema Inference with safe type coercion, length introspection & metadata tags
• ⚙️ Auto-Parameter Advisor for optimal salt count and shuffle parallelism
• 📉 Z-Score Plots and partition size diagnostics for visibility
• ✅ Fully PySpark-native — No RDDs, no CLI dependencies, no black-box wrappers

---

🧠 Quickstart

from hexadruid import HexaDruid

hd = HexaDruid(df)

# Step 1: Apply smart salting to balance skew
df_salted = hd.apply_smart_salting("sales_amount")

# Step 2 (Optional): Detect candidate primary or composite keys
key_info = hd.detect_keys()

# Step 3: Run schema optimizer + DRTree analyzer
typed_df, inferred_schema, dr_tree = HexaDruid.schemaVisor(df)

---

📚 What Does It Do?

Imagine a typical DataFrame:

order_id (UUID)	amount
a12e...	500.0
b98c...	5000.0
...	...

You're doing:

df.groupBy("amount").agg(...)

But most rows have the same amount, so Spark sends 99% of the work to 1 executor = skew 💥

---

🛠️ Main Classes & API Reference

Here are the key classes exposed by the HexaDruid package:

Class	Description
HexaDruid	Main entrypoint; handles smart salting, key detection, and schema optimization
SkewFeatureDetector	Detects skewed numeric columns for rebalancing
KeyFeatureDetector	Detects primary/composite keys (unique columns/combinations)
DRTree	Decision Rule Tree for logical sharding
AutoParameterAdvisor	Recommends optimal columns for skew balancing and groupBy
AdaptiveShuffleTuner	Tunes shuffle partition count dynamically
Branch, Root	(Advanced) Internal tree structure helpers

---

---

📑 API Reference

All core APIs are PySpark DataFrame-native. Below are the main classes and methods:

---

HexaDruid

Method	Description
HexaDruid(df, salt_count=10, output_dir="hexa_druid_outputs")	Initialize with a Spark DataFrame. Optional: set default salt bucket count and output directory.
apply_smart_salting(col_name, visualize=True, fine_tune=True, auto_tune=True, target_rows=1_000_000)	Balance skew in col_name by bucketizing (salting), visualizing, tuning salt count, and auto-tuning partitions. Returns a salted DataFrame with new columns: salt, salted_key.
detect_keys(dr_tree=None, composite_max_size=3, composite_threshold=0.99, verbose=True)	Detects the best candidate for a primary key or a composite key. Uses shard-aware logic if a dr_tree is passed. Returns a dict: {type, columns, confidence} or None.
repartition_on_salt(num_partitions=10)	Repartitions the salted DataFrame evenly on the salted_key column. Returns a repartitioned DataFrame.
show_partition_sizes(df, label="")	Prints the record count of each partition in the DataFrame, labeled for diagnostics.
build_shard_tree(detector, max_depth=3, min_samples=500)	Recursively build a logical sharding tree (DRTree) by splitting on skewed columns. Detector must implement .detect().
analyze_distribution(col_name)	Returns distribution stats (p95, p05, mean, std) for a given column.
schemaVisor(df, sample_fraction=0.2, max_depth=3, min_samples=500, skew_thresh=0.1, skew_top_n=3) (static)	Infers schema types, casts columns, and builds a DRTree for sharding. Returns a tuple: (typed_df, StructType, DRTree).
infer_numeric_columns(df) (static)	Returns a list of numeric columns (int, float, double, long, bigint).
detect_low_cardinality_categorical(df) (static)	Finds the first string column with ≤20 unique values (good for groupBy). Raises ValueError if none found.
timeit(func, label="") (static)	Times the execution of a function and logs the result.
_plot_comparison(col_name, df2) (private)	Generates and saves z-score barplots for original vs salted columns.

---

SkewFeatureDetector

Method	Description
SkewFeatureDetector(threshold=0.1, top_n=3)	Initialize detector with skew threshold and number of columns to return.
detect(df)	Returns the top-N most skewed numeric columns by quartile-based score.

---

KeyFeatureDetector

Method	Description
KeyFeatureDetector(verbose=False)	Initialize detector for verbose logging.
detect(df, dr_tree=None)	Returns best candidate key columns (single/composite), optionally shard-aware. List[str].
detectPrimaryKey(df, dr_tree=None, confidence_threshold=0.99, verbose=False) (static)	Detects primary key with confidence score. Returns dict or None.
detectCompositeKey(df, dr_tree=None, max_combination_size=3, confidence_threshold=0.99, verbose=False) (static)	Detects composite keys by testing combinations of 2-3 columns. Returns dict or None.

---

AutoParameterAdvisor

Method	Description
AutoParameterAdvisor(df, skew_top_n=3, cat_top_n=3)	Initialize with DataFrame and how many top columns to suggest.
recommend()	Returns (skew candidates, groupBy candidates, metrics DataFrame).
advise()	Interactive prompt: pick skew/groupBy columns (returns 2 strings).

---

DRTree

Method	Description
DRTree()	Create a new, empty Decision Rule Tree.
add_root(root)	Add a Root node to the DRTree.
to_dict()	Returns a JSON-serializable dictionary of the tree.

---

Branch, Root, DecisionNode, LeafNode

• Branch: Simple data holder for a leaf/shard predicate and its name.
• Root: Logical tree root node; holds one or more branches.
• DecisionNode: Internal tree node representing a split on a numeric column.
• LeafNode: Terminal node; represents a filtered logical subset ("shard") of your data.

---

balance_skew

Function	Description
balance_skew(df, output_dir="hexa_druid_outputs", partitions=10, verbose=False)	Runs full salting pipeline interactively: prompts for skew/groupBy columns, applies salting, shows before/after partition diagnostics. Returns a new DataFrame.

---

AdaptiveShuffleTuner

Method	Description
tune(spark, df, target_rows=1_000_000) (static)	Auto-tunes shuffle partitions based on target rows per partition. Returns repartitioned DataFrame.

---

Tip:

• See the tests/ directory for working code samples and usage patterns.
• All methods are intended for use with Spark DataFrames (PySpark >= 3.5.1).
• Advanced users can directly use the KeyFeatureDetector, SkewFeatureDetector, and AutoParameterAdvisor in custom pipelines.

---

⚖️ Smart Salting to the Rescue

df2 = hd.apply_smart_salting("amount")

What happens?

 Step 1: Analyze column distribution via IQR or Z-score
 Step 2: Generate N percentile buckets
 Step 3: Assign salt ID per row using bucket bounds
 Step 4: Create salted_key = amount_salt
 Step 5: Repartition on salted_key for parallelism

📈 This rebalances the shuffle phase for joins, groupBy, and aggregates.

---

🧠 DRTree Explained Visually

The DRTree is a decision-rule tree, not a classifier.

It recursively splits data into shards by applying SQL-style predicates. Each leaf is a filtered logical subset of the DataFrame.

                        [Root: sales_amount]
                                |
                   ┌───────────┴────────────┐
        [amount <= 500]             [amount > 500]
               |                           |
       ┌───────┴───────┐            ┌──────┴───────┐
 [amount <= 100] [>100, ≤500]   [>500, ≤1000]  [>1000]
       |         |                  |             |
   [Leaf A]   [Leaf B]          [Leaf C]       [Leaf D]
 (shard_1)   (shard_2)         (shard_3)      (shard_4)

Each leaf holds:

• Filtered subset of the DataFrame (as a Spark SQL query)
• Associated metadata like row count, min/max, schema drift
• Auto key detection can run within these shards

---

🔬 Leaf-Level Parallelization

DRTree enables parallel insight:

• Each leaf is autonomous (you can infer schema, key, and stats per leaf)
• Makes the system robust to changes over time (drift detection)
• Enables controlled analytics:

[DRTree Output]
Leaf A:
  - rows: 30K
  - key_confidence: 0.92
  - type: Float(5,2)

Leaf D:
  - rows: 300K (hotspot!)
  - key_confidence: 0.12
  - type: String(255)

---

🔑 Key Detection (Optional & Shard-Aware)

key_info = hd.detect_keys()

You don’t need to force primary keys.

This is just analysis — it evaluates uniqueness confidence for each column (or combination of columns):

• Primary Key:

score = (approx_count_distinct(col) / total_rows) - null_ratio

If score ≥ 0.99, it’s a good candidate.

• Composite Key:

combo_key = concat_ws("_", col1, col2, ...)
score = approx_count_distinct(combo_key) / total_rows - null_ratio

DRTree passes its shard filters into detect_keys() to evaluate keys per subgroup — boosting accuracy.

---

🧠 Smart Salting Internals

🧪 Step-by-step:

1. Detect Skew

   • If z_score range is too large or
   • IQR shows asymmetry (Q3 - Q2 ≫ Q2 - Q1)

2. Split by Percentiles

percentiles = percentile_approx("amount", [0.0, 0.1, ..., 1.0])

3. Salt Bucketing Logic

salt = when(col >= p0 & col < p1, 0) \
     .when(col >= p1 & col < p2, 1) ...

4. Create Salted Key

salted_key = concat_ws("_", col("amount"), col("salt"))
df = df.withColumn("salted_key", salted_key).repartition("salted_key")

5. Auto-Tune Salt Count

• If distribution is dense, fewer buckets suffice
• Otherwise, more salting is applied dynamically

---

📈 Visualization Example

Output from schemaVisor():

Leaf Node A [shard_0]
- size: 102,391
- type: Float(8,2)
- confidence: 92%

Leaf Node B [shard_1]
- size: 489,128 (dense zone)
- skew detected!
- Recommended salt count: 10

You can visualize the Z-score distribution:

Before:
  [■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■             ]

After:
  [■■■■■■■■■■■■■■■■■■■■■      ■■■■■■■■■■■■■■■■■■■■■■   ]

---

🧪 Testing

pytest tests/

Mocked SparkSession with synthetic data is used to ensure full coverage.

---

🧱 Suggested Project Structure

hexadruid/
├── __init__.py
├── core.py                # HexaDruid entry point
├── skew_balancer.py       # Smart salting logic
├── drtree.py              # DRTree shard splitting
├── key_detection.py       # Unique key checker
├── schema_optimizer.py    # Type inference
├── advisor.py             # Parameter tuning
├── utils.py               # Logging, plots, etc.
└── tests/                 # Test suite

---

🔧 Roadmap

• CLI interface
• Delta Lake + Iceberg support
• JupyterLab extension
• DRTree JSON export for audits
• Cost metrics estimation
• Column statistics and visualization dashboard

---

📄 License

MIT License

---

🤝 Contributing

Pull requests, ideas, and contributions welcome!

We believe Spark shouldn’t be slow. Let’s make it smarter together.

---