Back to Advanced Data Engineering
Q18Question 18 of 24

GraphX and MLlib

Key Concepts
  • GraphX models a graph as distributed vertex and edge RDDs and provides algorithms like PageRank, connected components, and triangle counting.
  • MLlib provides distributed machine learning: classification, regression, clustering, and recommendation that scale across a cluster.
  • The modern MLlib API is the DataFrame-based spark.ml Pipeline that chains transformers and estimators into one reproducible workflow.
  • GraphX is RDD-based and somewhat legacy; GraphFrames is preferred for new graph work with DataFrame queries.
  • GraphX and MLlib combine well: graph features (PageRank, component size) become inputs to an MLlib fraud classifier.
  • Spark MLlib targets classical algorithms at scale; deep learning is handled by dedicated frameworks integrated with Spark.

Definition

GraphX does graph-parallel computation on vertices and edges (PageRank, connected components, triangle counting); MLlib does distributed machine learning (classification, regression, clustering, recommendation) with a Pipeline API.
  • GraphX – Spark's library for graph-parallel computation, representing a graph as distributed vertex and edge RDDs, and providing algorithms such as PageRank, connected components, and triangle counting.
  • MLlib – Spark's distributed machine learning library, offering classification, regression, clustering, and recommendation, plus feature engineering through a DataFrame-based Pipeline API.
  • Pipeline API – chains feature transformers (StringIndexer, VectorAssembler) and estimators into a single fit-and-transform workflow for reproducible ML.
  • Combination – graph features from GraphX can feed into MLlib classifiers, making them powerful together for fraud detection.

At the National Bank of India, GraphX is the tool for questions about relationships – who sends money to whom. Draw every account as a dot and every transfer as an arrow, and GraphX can find rings of accounts that move money among themselves (connected components), rank which accounts are most central to the money flow (PageRank), and spot tight triangles of mutual transfers that may signal collusion. MLlib is the tool for learning from data – training a model on past transactions to predict which new ones are likely fraudulent, grouping customers by behavior, or recommending products – and its Pipeline API is a fixed assembly line that cleans, prepares, and runs the model the same way every time.

Understanding the Concepts

What is GraphX?

GraphX models a graph as two distributed collections: a vertex RDD of node identifiers with attributes, and an edge RDD of source-to-destination pairs with attributes.

PYTHON
# GraphX is Scala-only; here is the equivalent using GraphFrames in PySpark
from graphframes import GraphFrame
from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("GraphDemo").getOrCreate()

# Vertices: bank accounts
vertices = spark.createDataFrame([
    ("ACC-1001", "Rajesh", "Mumbai"),
    ("ACC-1002", "Priya",  "Delhi"),
    ("ACC-1003", "Pankaj", "Chennai"),
], ["id", "holder_name", "branch"])

# Edges: money transfers
edges = spark.createDataFrame([
    ("ACC-1001", "ACC-1002", 50000),
    ("ACC-1002", "ACC-1003", 30000),
    ("ACC-1003", "ACC-1001", 20000),  # circular transfer
], ["src", "dst", "amount"])

g = GraphFrame(vertices, edges)

The graph abstraction supports algorithms that operate on the structure of relationships.

What algorithms does GraphX provide?

GraphX ships with several built-in algorithms that answer different structural questions.

PYTHON
# PageRank: which accounts are most central to the money flow?
ranks = g.pageRank(resetProbability=0.15, maxIter=10)
ranks.vertices.select("id", "holder_name", "pagerank").show()
TEXT
+--------+-----------+------------------+
|      id|holder_name|          pagerank|
+--------+-----------+------------------+
|ACC-1001|     Rajesh|1.0654321098765432|
|ACC-1002|      Priya|0.9672839506172839|
|ACC-1003|     Pankaj|0.9672839506172839|
+--------+-----------+------------------+
PYTHON
# Connected Components: find clusters of linked accounts
components = g.connectedComponents()
components.select("id", "component").show()

# Triangle Counting: measure local connectivity
triangles = g.triangleCount()
triangles.select("id", "count").show()

What is MLlib?

MLlib provides distributed machine learning algorithms that train on datasets far larger than one machine's memory.

PYTHON
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import VectorAssembler, StringIndexer
from pyspark.sql import functions as F

# Prepare training data for fraud detection
txn = spark.read.parquet("/lake/silver/transactions")
txn = txn.withColumn("is_fraud", F.col("fraud_flag").cast("double"))

assembler = VectorAssembler(
    inputCols=["amount", "tx_count", "avg_balance"],
    outputCol="features"
)
train_data = assembler.transform(txn)

# Train a logistic regression classifier
lr = LogisticRegression(featuresCol="features", labelCol="is_fraud")
model = lr.fit(train_data)

# Predict on new transactions
predictions = model.transform(train_data)
predictions.select("account_id", "is_fraud", "prediction").show(5)

What is the Pipeline API?

The Pipeline API chains transformers and estimators into a single reproducible workflow, preventing training-serving skew.

PYTHON
from pyspark.ml import Pipeline
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.feature import StringIndexer, VectorAssembler

# Define the pipeline stages
channel_indexer = StringIndexer(inputCol="channel", outputCol="channel_idx")
assembler = VectorAssembler(
    inputCols=["amount", "tx_count", "avg_balance", "channel_idx"],
    outputCol="features"
)
rf = RandomForestClassifier(
    featuresCol="features", labelCol="is_fraud", numTrees=100
)

pipeline = Pipeline(stages=[channel_indexer, assembler, rf])

# fit() learns statistics and trains the model
pipeline_model = pipeline.fit(train_data)

# transform() applies the same preprocessing + model to new data
predictions = pipeline_model.transform(new_txn)

The pipeline ensures the exact same preprocessing runs during training and inference.

How do GraphX and MLlib combine for fraud detection?

Graph features from GraphX become inputs to an MLlib classifier, creating a powerful fraud detection system.

PYTHON
# Step 1: Extract graph features
ranks = g.pageRank(resetProbability=0.15, maxIter=10).vertices
components = g.connectedComponents()

graph_features = ranks.join(components, "id") \
    .select(
        F.col("id").alias("account_id"),
        F.col("pagerank"),
        F.col("component").alias("component_id")
    )

# Step 2: Join graph features with transaction features
enriched = txn.join(graph_features, "account_id")

# Step 3: Build pipeline with graph + transaction features
assembler = VectorAssembler(
    inputCols=["amount", "tx_count", "avg_balance", "pagerank"],
    outputCol="features"
)
pipeline = Pipeline(stages=[assembler, rf])
model = pipeline.fit(enriched)

A high-value transaction from a central node inside a tight ring scores much higher than a normal transaction from an isolated pair.

What ML algorithm families does MLlib support?

MLlib covers the standard families of machine learning.

| Family | Algorithms | Banking Use Case |
|---|---|---|
| Classification | Logistic regression, decision trees, random forest, GBT | Fraud detection |
| Regression | Linear regression, decision tree, GBT | Credit scoring, amount prediction |
| Clustering | K-means, Gaussian mixture, LDA | Customer segmentation |
| Recommendation | ALS (collaborative filtering) | Product recommendation |
| Feature Engineering | StringIndexer, OneHotEncoder, VectorAssembler, scalers | Pipeline preprocessing |

PYTHON
from pyspark.ml.clustering import KMeans

kmeans = KMeans(k=5, seed=42, featuresCol="features")
clusters = kmeans.fit(customer_features)
clusters.transform(customer_features).select("account_id", "prediction").show()

What are the ecosystem directions and limitations?

GraphX is RDD-based and somewhat legacy. GraphFrames, built on DataFrames, is preferred for new graph work.

PYTHON
# GraphX: Scala-only, RDD-based, legacy
# GraphFrames: PySpark-friendly, DataFrame-based, preferred
from graphframes import GraphFrame

# MLlib RDD API: deprecated
# spark.ml (DataFrame API): current and recommended
from pyspark.ml.classification import LogisticRegression  # spark.ml

Spark MLlib targets large-scale distributed training of classical algorithms. Deep learning is handled by dedicated frameworks (TensorFlow, PyTorch) integrated with Spark through libraries like Horovod.

What should a senior engineer know?

  • Use spark.ml (DataFrame-based) for all new ML work; the RDD-based spark.mllib is deprecated.
  • Use GraphFrames for new graph work; GraphX is Scala-only and RDD-based.
  • Pipelines make ML reproducible by bundling preprocessing and model into one serializable object.
  • Spark MLlib excels at distributed classical ML. For deep learning, integrate a dedicated framework.
  • Graph + ML combination is a powerful pattern for fraud, AML, and network analysis in banking.
PYTHON
# Save and load the pipeline model for production serving
pipeline_model.save("/models/fraud_pipeline_v1")
loaded = PipelineModel.load("/models/fraud_pipeline_v1")

What are common interview follow-up questions?

  • Can you do deep learning with MLlib? MLlib does not have built-in deep learning. Use TensorFlow or PyTorch with Spark for data loading.
  • What is training-serving skew? When preprocessing differs between training and inference, causing wrong predictions. Pipelines prevent this.
  • Can GraphX handle billions of edges? Yes, but performance depends on the cluster size and the algorithm's communication pattern.

Visual Overview

DIAGRAM
  GraphX + MLlib Fraud Detection Pipeline
  ────────────────────────────────────────

  ┌─────────────────────┐
  │  Account Transfers   │
  │  (vertices + edges) │
  └──────────┬──────────┘
             │
     GraphX / GraphFrames
             │
  ┌──────────▼──────────┐
  │  Graph Features      │
  │  - PageRank          │
  │  - Component ID      │
  │  - Triangle Count    │
  └──────────┬──────────┘
             │  join
  ┌──────────▼──────────┐
  │  Transaction Features│
  │  - amount            │
  │  - tx_count          │
  │  - avg_balance       │
  └──────────┬──────────┘
             │
     MLlib Pipeline
     (StringIndexer → VectorAssembler → Classifier)
             │
  ┌──────────▼──────────┐
  │  Fraud Predictions   │
  └─────────────────────┘

Key Takeaway

GraphX is Spark's library for graph-parallel computation, representing a graph as distributed vertex and edge collections. It provides algorithms like PageRank to rank vertices by centrality, connected components to find clusters of linked nodes, and triangle counting to measure local connectivity. In banking, these answer questions about who sends money to whom, which accounts are central to the flow, and where tight rings of mutual transfers might signal collusion.

MLlib is Spark's distributed machine learning library, covering classification, regression, clustering, and recommendation with algorithms that train across a cluster. Its modern core is the DataFrame-based Pipeline API, which chains feature transformers like StringIndexer and VectorAssembler with an estimator into one reproducible workflow. Fitting the pipeline learns statistics and trains the model; transforming applies the entire sequence to new data, ensuring the same preprocessing runs during training and inference.

At the National Bank of India these worked together in fraud detection. We used connected components on the transfer graph to find rings of collusive accounts and PageRank to identify central nodes. The component size and PageRank score became features fed into an MLlib classification pipeline alongside transaction features like amount and frequency. A high-value transaction from a central node inside a tight ring scored much higher than a normal transaction from an isolated pair.

The ecosystem nuances to know are that GraphX is RDD-based and Scala-only, with GraphFrames preferred for new work, that the RDD-based spark.mllib is deprecated in favor of spark.ml, and that Spark MLlib targets classical algorithms at scale, not deep learning.

Remember one sentence: GraphX answers relationship questions with PageRank, connected components, and triangle counting, while MLlib builds reproducible distributed ML with the Pipeline API, and combining graph features with transaction features creates powerful fraud classifiers.

---

Share
Use ← → arrow keys to navigate