Skip to content
Tutorial Reference

Python PySpark: How to Slice a PySpark DataFrame into Two Row-Wise DataFrames

How to split a PySpark DataFrame into two row-wise DataFrames using randomSplit(), limit(), and other approaches, with tips for handling large datasets, duplicates, and memory usage.

Splitting a PySpark DataFrame into two smaller DataFrames by rows is a common operation in data processing - whether you need to create training and test sets, separate data for parallel processing, or paginate large datasets. Unlike Pandas, PySpark DataFrames don't support direct row indexing, so slicing requires different approaches.

This guide covers four methods to split a PySpark DataFrame row-wise, each suited to different use cases.

Sample DataFrame

All examples use the following PySpark DataFrame:

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName('SlicingDemo').getOrCreate()

rows = [
    ['James Anderson', 69, 'New York'],
    ['Emily Chen', 66, 'Chicago'],
    ['Michael Davis', 9, 'Houston'],
    ['Sarah Martinez', 15, 'Phoenix']
]

columns = ['Player', 'Titles', 'City']
df = spark.createDataFrame(rows, columns)
df.show()

Output:

+--------------+------+--------+
|        Player|Titles|    City|
+--------------+------+--------+
|James Anderson|    69|New York|
|    Emily Chen|    66| Chicago|
| Michael Davis|     9| Houston|
|Sarah Martinez|    15| Phoenix|
+--------------+------+--------+

Method 1: Using limit() and subtract()

This approach takes the first N rows with limit() and gets the remaining rows by subtracting the first slice from the original DataFrame.

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName('SlicingDemo').getOrCreate()

rows = [
    ['James Anderson', 69, 'New York'],
    ['Emily Chen', 66, 'Chicago'],
    ['Michael Davis', 9, 'Houston'],
    ['Sarah Martinez', 15, 'Phoenix']
]

columns = ['Player', 'Titles', 'City']
df = spark.createDataFrame(rows, columns)

df1 = df.limit(3)

df2 = df.subtract(df1)

print("First slice:")
df1.show()

print("Second slice:")
df2.show()

Output:

First slice:
+--------------+------+--------+
|        Player|Titles|    City|
+--------------+------+--------+
|James Anderson|    69|New York|
|    Emily Chen|    66| Chicago|
| Michael Davis|     9| Houston|
+--------------+------+--------+

Second slice:
+--------------+------+-------+
|        Player|Titles|   City|
+--------------+------+-------+
|Sarah Martinez|    15|Phoenix|
+--------------+------+-------+

Method 2: Using randomSplit()

The randomSplit() method splits a DataFrame into multiple parts based on specified weight ratios. Rows are distributed randomly, making this ideal for creating training and test datasets.

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName('SlicingDemo').getOrCreate()

rows = [
    ['James Anderson', 69, 'New York'],
    ['Emily Chen', 66, 'Chicago'],
    ['Michael Davis', 9, 'Houston'],
    ['Sarah Martinez', 15, 'Phoenix']
]

columns = ['Player', 'Titles', 'City']
df = spark.createDataFrame(rows, columns)

df1, df2 = df.randomSplit([0.7, 0.3], seed=42)

print(f"First slice ({df1.count()} rows):")
df1.show()

print(f"Second slice ({df2.count()} rows):")
df2.show()

Output (varies by run unless seed is set):

First slice (3 rows):
+--------------+------+-------+
|        Player|Titles|   City|
+--------------+------+-------+
|    Emily Chen|    66|Chicago|
| Michael Davis|     9|Houston|
|Sarah Martinez|    15|Phoenix|
+--------------+------+-------+

Second slice (1 rows):
+--------------+------+--------+
|        Player|Titles|    City|
+--------------+------+--------+
|James Anderson|    69|New York|
+--------------+------+--------+

Method 3: Using collect() and List Slicing

This method converts the DataFrame to a list of Row objects, slices the list using Python's standard slicing, and converts the slices back to PySpark DataFrames.

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName('SlicingDemo').getOrCreate()

rows = [
    ['James Anderson', 69, 'New York'],
    ['Emily Chen', 66, 'Chicago'],
    ['Michael Davis', 9, 'Houston'],
    ['Sarah Martinez', 15, 'Phoenix']
]

columns = ['Player', 'Titles', 'City']
df = spark.createDataFrame(rows, columns)

row_list = df.collect()

part1 = row_list[:1]   # First row
part2 = row_list[1:]   # Remaining rows

slice1 = spark.createDataFrame(part1, columns)
slice2 = spark.createDataFrame(part2, columns)

print("First slice:")
slice1.show()

print("Second slice:")
slice2.show()

Output:

First slice:
+--------------+------+--------+
|        Player|Titles|    City|
+--------------+------+--------+
|James Anderson|    69|New York|
+--------------+------+--------+

Second slice:
+--------------+------+-------+
|        Player|Titles|   City|
+--------------+------+-------+
|    Emily Chen|    66|Chicago|
| Michael Davis|     9|Houston|
|Sarah Martinez|    15|Phoenix|
+--------------+------+-------+

This approach gives you exact control over which rows go into each slice based on their position.

Method 4: Converting to Pandas and Using iloc[]

This method converts the PySpark DataFrame to a Pandas DataFrame, uses Pandas' iloc[] for precise row slicing, and converts the results back to PySpark.

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName('SlicingDemo').getOrCreate()

rows = [
    ['James Anderson', 69, 'New York'],
    ['Emily Chen', 66, 'Chicago'],
    ['Michael Davis', 9, 'Houston'],
    ['Sarah Martinez', 15, 'Phoenix']
]

columns = ['Player', 'Titles', 'City']
df = spark.createDataFrame(rows, columns)

pandas_df = df.toPandas()

pdf1 = pandas_df.iloc[:2]   # First 2 rows
pdf2 = pandas_df.iloc[2:]   # Last 2 rows

df1 = spark.createDataFrame(pdf1)
df2 = spark.createDataFrame(pdf2)

print("First slice:")
df1.show()

print("Second slice:")
df2.show()

Output:

First slice:
+--------------+------+--------+
|        Player|Titles|    City|
+--------------+------+--------+
|James Anderson|    69|New York|
|    Emily Chen|    66| Chicago|
+--------------+------+--------+

Second slice:
+--------------+------+-------+
|        Player|Titles|   City|
+--------------+------+-------+
| Michael Davis|     9|Houston|
|Sarah Martinez|    15|Phoenix|
+--------------+------+-------+

This gives you the full power of Pandas slicing, including step values (iloc[::2]), negative indexing, and more.

Method Comparison

Method Preserves Order Exact Split Handles Duplicates Large Data Safe
limit() + subtract() No Yes (by count) No (removes duplicates) Yes
randomSplit() Yes Approximate Yes Yes
collect() + list slicing Yes Yes (by position) Yes No (memory risk)
toPandas() + iloc[] Yes Yes (by position) Yes No (memory risk)

When to Use Each Method

  • randomSplit(): Best for machine learning train/test splits where random distribution matters and exact counts are not critical.
  • limit() + subtract(): Best for distributed environments where you need the first N rows without collecting data to the driver. Be aware of duplicate handling.
  • collect() or toPandas(): Best for small DataFrames where you need precise positional slicing and full control over the split point.

Choose the method that fits your data size and whether you need deterministic, position-based slicing or approximate, distributed splitting.

Navigation