Using Arrow for ML

---

The combination of Arrow and Spark

Essentially, the "Marriage" works because Spark (which runs on the JVM/Java) and Python (which runs on CPython) speak different languages. Without Arrow, they have to "pack and unpack" every single row as they talk to each other.

The "Why": The Python-Spark Bottleneck

Before Arrow integration, when you used a Pandas UDF (User-Defined Function) in Spark, the following happened:

1. Spark (Java) would take its internal data.

2. It would serialize it into a generic row-based format (Py4J).

3. Python would receive those rows and re-serialize them into a Pandas DataFrame.

4. After processing, it would do the whole thing in reverse to get the data back to Spark.

This "translation" was often slower than the actual data processing!

The Arrow Solution: Shared Memory

When you enable Arrow in Spark, the data isn't "translated" row-by-row. Instead:

1. Spark organizes the data into Arrow RecordBatches (the columnar format we've been discussing).

2. It hands those batches directly to Python.

3. Because Pandas (and NumPy) can "read" Arrow memory natively, Python starts working instantly with zero-copy or near-zero-copy overhead.

When NOT to use it

The book is right to warn you. If your entire pipeline is just Spark SQL (e.g., df.select("col").groupBy(...)), you should not force it into Arrow.

• Spark's Native Format: Optimized for Java/Scala operations and shuffle logic.

• The Cost: Converting Spark's internal "UnsafeRow" format to Arrow takes CPU cycles. If you don't need to leave the JVM, don't pay that tax.

---

Example: Spark + Arrow: Optimized Data Transfer

When you use the createDataFrame method with Arrow enabled, Spark bypasses the row-by-row serialization. It converts the Pandas DataFrame into an Arrow Table and streams the record batches directly into the JVM.

import pyarrow.csv as pv
import pyarrow as pa
from pyspark.sql import SparkSession

# Enable Arrow optimization
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

# 1. Read CSV using PyArrow (High-speed C++ engine)
# 2. Convert to Pandas
# 3. Create Spark DataFrame (Zero-copy optimization triggered here)
pdf = pv.read_csv('large_dataset.csv').to_pandas()
df = spark.createDataFrame(pdf)

df.show()

---

Example: Pandas UDFs (Vectorized Operations)

The "Standard" Spark UDF is slow because it executes your Python code one row at a time. A Pandas UDF (also called a Vectorized UDF) sends a whole Arrow Record Batch to Python as a Series or DataFrame, allowing you to use optimized libraries like NumPy.

Example: Vectorized Math with NumPy

import pandas as pd
import numpy as np
from pyspark.sql.functions import pandas_udf

# Define the UDF with a type hint (Series to Series)
@pandas_udf("double")
def vectorized_math(v: pd.Series) -> pd.Series:
    # This runs on a whole batch of data at once using NumPy
    return np.log(v + 1) * 10 

# Apply it to your Spark DataFrame
df_with_calc = df.withColumn("processed_val", vectorized_math(df["col_name"]))
df_with_calc.show()

Another example on UDFs in Spark

We will generate some synthetic taxi data, read it using the high-speed PyArrow engine, move it into Spark, and then apply the Grouped Map (Vectorized UDF) to normalize the fares.

This example demonstrates the "Marriage" of the PyArrow C++ Reader and the Spark Distributed Engine.

The Full Pipeline

import pyarrow as pa
import pyarrow.csv as pv
import pandas as pd
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, DoubleType, LongType, StringType
import pyspark.sql.functions as F

# 1. Initialize Spark with Arrow Optimization Enabled
spark = SparkSession.builder \
    .appName("ArrowSparkMarriage") \
    .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
    .getOrCreate()

# 2. Setup: Create a dummy CSV to simulate your NYC Taxi data
data = {
    'VendorID': [1, 1, 1, 2, 2, 2],
    'pickup_month': ['Jan', 'Jan', 'Jan', 'Jan', 'Jan', 'Jan'],
    'total_amount': [10.0, 20.0, 30.0, 5.0, 15.0, 25.0]
}
pd.DataFrame(data).to_csv('taxi_data.csv', index=False)

# 3. Best Practice Read: Use PyArrow to read, then transfer to Spark
# This triggers the 'ParallelCollectionRDD' optimization
arrow_table = pv.read_csv('taxi_data.csv')
df = spark.createDataFrame(arrow_table.to_pandas())

# 4. Define the output Schema for the UDF
# We are keeping original columns and adding 'normalized_fare'
output_schema = StructType(df.schema.fields + [
    StructField('normalized_fare', DoubleType())
])

# 5. Define the Vectorized (Pandas) UDF Logic
def vector_normalize(pdf):
    """
    pdf is a pandas.DataFrame representing a group
    """
    fare = pdf.total_amount
    # Perform vectorized math using NumPy/Pandas logic
    pdf['normalized_fare'] = (fare - fare.mean()) / fare.std()
    return pdf

# 6. Execute the Grouped Map operation
# This distributes the groups across the cluster and uses Arrow IPC for speed
df_normalized = df.groupby('VendorID', 'pickup_month') \
                  .applyInPandas(vector_normalize, schema=output_schema)

# 7. Show the Result
df_normalized.select('VendorID', 'total_amount', 'normalized_fare').show()

---

What’s happening under the hood?

1. The Read: PyArrow parses the CSV using its multi-threaded C++ engine. This is significantly faster than Spark's native CSV reader for local or single-node files.

2. The Transfer: Because spark.sql.execution.arrow.pyspark.enabled is true, the createDataFrame call doesn't pickle the data. It sends it as a raw binary stream of Arrow RecordBatches.

3. The Distribution: Spark identifies that the data needs to be grouped by VendorID. It shuffles the data so all Vendor 1 data is in one "bucket."

4. The Execution: Each "bucket" is converted back to an Arrow batch, sent to a Python worker, and turned into a Pandas DataFrame.

5. The Math: Your normalization logic runs. Notice how fare.mean() and fare.std() work on the whole group column at once.

6. The Return: The results stream back to the JVM via Arrow, and Spark re-assembles them into the final distributed DataFrame.

Use case context (for example for NYC Taxi data):

• Feature Engineering: Normalizing prices or distances per borough.

• Time-Series Analysis: Calculating rolling averages within a specific time window.

• Batch Inference: Running a local machine learning model against every taxi zone's data independently.

---

Handling Parquet

Best Practice: The "Small-to-Medium" Optimization

Use the PyArrow + createDataFrame approach when your data fits comfortably in the memory of your driver node, but you want Spark's distributed engine to handle the actual processing or ML training.

Code Example: PyArrow-First Reading

import pyarrow.parquet as pq
import pyarrow as pa
from pyspark.sql import SparkSession

# 1. Enable Arrow in Spark
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

# 2. Use PyArrow to read (Multi-threaded C++ engine)
# Best practice: Only read the columns you need to save memory
table = pq.read_table('my_data.parquet', columns=['id', 'target', 'feature_1'])

# 3. Hand off to Spark
# This creates a ParallelCollectionRDD
sdf = spark.createDataFrame(table.to_pandas())

# Spark now takes this 'collection' and distributes it across executors
# far more efficiently than if it had to manage the file I/O itself.
sdf.groupBy("target").count().show()

• Native Spark (FileScanRDD): Spark treats the filesystem as the source of truth. It looks at the file sizes and splits them into partitions based on maxPartitionBytes. This is usually the default for "Big Data" (TBs).

• PyArrow Ingestion (ParallelCollectionRDD): You use Python’s high-performance C++ engine to read a file into memory as an Arrow Table, then hand it to Spark. Spark sees this as an "in-memory collection" and parallelizes it across the cluster.

Why this works better (Technical Context)

Aspect	Spark Native Read	PyArrow + Spark
I/O Engine	Java-based (Hadoop FileInputFormat)	C++ based (Multi-threaded Arrow)
Parallelism	Based on file splits / HDFS blocks	Based on Record Batches in memory
RDD Type	FileScanRDD	ParallelCollectionRDD
Ideal For	Massive datasets (S3/HDFS)	Local files or Cloud files < 10GB

---

Machine Learning Use Case

This is where it becomes relevant to your ML reading. Most ML libraries (Scikit-learn, PyTorch, TensorFlow, XGBoost) are written in C++ or Python. They cannot read Spark's internal Java memory.

The ML Workflow:

1. Spark: Does the heavy lifting (joining TBs of data, cleaning).

2. Arrow: "Pipes" the cleaned data out of the JVM.

3. XGBoost/PyTorch: Receives the Arrow batches and starts training immediately.

---