How to find the size or shape of a DataFrame in PySpark?

Question

I am trying to find out the size/shape of a DataFrame in PySpark. I do not see a single function that can do this.

In Python, I can do this:

data.shape()

Is there a similar function in PySpark? This is my current solution, but I am looking for an element one

row_number = data.count()
column_number = len(data.dtypes)

The computation of the number of columns is not ideal...

You mean `data.shape` for NumPy and Pandas? `shape` is not a function. — flow2k, Jul 03 '19 at 08:44
What is not ideal? I am not sure what else you would like to accomplish than what you already have (except for replacing data.dtypes with data.columns, but it makes little difference). — Melkor.cz, Aug 27 '20 at 11:55

score 253 · Answer 1 · edited Sep 02 '20 at 14:22

253

You can get its shape with:

print((df.count(), len(df.columns)))

edited Sep 02 '20 at 14:22

yatu

answered Aug 11 '17 at 17:28

George Fisher

Will this work fine for larger datasets spread across nodes? – Purushothaman Srikanth Jul 07 '21 at 07:45
2

Why doesn't Pyspark Dataframe simply store the shape values like pandas dataframe does with `.shape`? Having to call count seems incredibly resource-intensive for such a common and simple operation. – THIS USER NEEDS HELP Feb 24 '22 at 04:06

score 80 · Answer 2 · edited Jan 08 '19 at 06:32

80

Use df.count() to get the number of rows.

edited Jan 08 '19 at 06:32

Tshilidzi Mudau

answered Aug 18 '17 at 13:33

VMEscoli

Louis Yang · Answer 3 · 2020-11-08T00:04:13.183

55

Add this to the your code:

import pyspark
def spark_shape(self):
    return (self.count(), len(self.columns))
pyspark.sql.dataframe.DataFrame.shape = spark_shape

Then you can do

>>> df.shape()
(10000, 10)

But just remind you that .count() can be very slow for very large table that has not been persisted.

edited Nov 08 '20 at 00:04

answered Dec 19 '18 at 19:20

Louis Yang

5

I really think it's a bad idea to change the DataFrame API, without a valid reason to do so. just call `spark_shape(my_df)`... Moreover, possibly name the function something clearer like `compute_dataframe_shape`... – Davide Aug 10 '22 at 13:50
This tip is for data scientists or data analysts who have to constantly type those command every day multiple times when they do analysis on data. It is not for better engineering perspective or production code. – Louis Yang Sep 05 '22 at 09:07

score 12 · Answer 4 · answered Mar 31 '20 at 11:51

12

print((df.count(), len(df.columns)))

is easier for smaller datasets.

However if the dataset is huge, an alternative approach would be to use pandas and arrows to convert the dataframe to pandas df and call shape

spark.conf.set("spark.sql.execution.arrow.enabled", "true")
spark.conf.set("spark.sql.crossJoin.enabled", "true")
print(df.toPandas().shape)

answered Mar 31 '20 at 11:51

Venzu251720

17

Isn't .toPandas an action? Meaning: isn't this going to collect the data to your master, and then call shape on it? If so, it would be inadvisable to do that, unless you're sure it's going to fit in master's memory. – ponadto Apr 06 '20 at 16:57
8

If the dataset is huge, collecting to Pandas is exactly what you do NOT want to do. Btw: Why do you enable cross join for this? And does the arrow configuration help collecting to pandas? – Melkor.cz Aug 27 '20 at 11:49

score 4 · Answer 5 · answered Oct 30 '16 at 13:34

4

I think there is not similar function like data.shape in Spark. But I will use len(data.columns) rather than len(data.dtypes)

answered Oct 30 '16 at 13:34

YungChun

8

that just gives you number of columns. What about number of rows? – JanLauGe Jul 23 '17 at 12:56

5 Answers5