I have a dataframe stage_name_valid_df which looks like follows on calling .show()
and .toPandas() works as well
but on calling stage_name_valid_df.count() I get an error as shown below. What could be the possible cause of it? My java version is openjdk version "1.8.0_342" and I am running pyspark locally on my laptop. Pyspark version is 3.3.0 and python version is 3.9.12.
java.lang.IllegalStateException: Couldn't find revision_id#17L in [stage_name#116,sum#1397L]
at org.apache.spark.sql.catalyst.expressions.BindReferences$$anonfun$bindReference$1.applyOrElse(BoundAttribute.scala:80)
at org.apache.spark.sql.catalyst.expressions.BindReferences$$anonfun$bindReference$1.applyOrElse(BoundAttribute.scala:73)
at org.apache.spark.sql.catalyst.trees.TreeNode.$anonfun$transformDownWithPruning$1(TreeNode.scala:584)
at org.apache.spark.sql.catalyst.trees.CurrentOrigin$.withOrigin(TreeNode.scala:176)
at org.apache.spark.sql.catalyst.trees.TreeNode.transformDownWithPruning(TreeNode.scala:584)
at org.apache.spark.sql.catalyst.trees.TreeNode.$anonfun$transformDownWithPruning$3(TreeNode.scala:589)
at scala.collection.immutable.List.map(List.scala:297)
at org.apache.spark.sql.catalyst.trees.TreeNode.mapChildren(TreeNode.scala:698)
at org.apache.spark.sql.catalyst.trees.TreeNode.transformDownWithPruning(TreeNode.scala:589)
at org.apache.spark.sql.catalyst.trees.TreeNode.transformDown(TreeNode.scala:560)
at org.apache.spark.sql.catalyst.trees.TreeNode.transform(TreeNode.scala:528)
at org.apache.spark.sql.catalyst.expressions.BindReferences$.bindReference(BoundAttribute.scala:73)
at org.apache.spark.sql.execution.GenerateExec.boundGenerator$lzycompute(GenerateExec.scala:75)
at org.apache.spark.sql.execution.GenerateExec.boundGenerator(GenerateExec.scala:75)
at org.apache.spark.sql.execution.GenerateExec.$anonfun$doExecute$3(GenerateExec.scala:95)
at scala.collection.Iterator$$anon$11.nextCur(Iterator.scala:486)
at scala.collection.Iterator$$anon$11.hasNext(Iterator.scala:492)
at scala.collection.Iterator$ConcatIterator.hasNext(Iterator.scala:224)
at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:460)
at org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage5.hashAgg_doAggregateWithoutKey_0$(Unknown Source)
at org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage5.processNext(Unknown Source)
at org.apache.spark.sql.execution.BufferedRowIterator.hasNext(BufferedRowIterator.java:43)
at org.apache.spark.sql.execution.WholeStageCodegenExec$$anon$1.hasNext(WholeStageCodegenExec.scala:760)
at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:460)
at org.apache.spark.shuffle.sort.BypassMergeSortShuffleWriter.write(BypassMergeSortShuffleWriter.java:140)
at org.apache.spark.shuffle.ShuffleWriteProcessor.write(ShuffleWriteProcessor.scala:59)
at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:99)
at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:52)
at org.apache.spark.scheduler.Task.run(Task.scala:136)
at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$3(Executor.scala:548)
at org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:1504)
at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:551)
at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
at java.lang.Thread.run(Thread.java:750)
A complete program to reproduce the error is below (although with a minimal toy dataset):
from pyspark.sql import SparkSession
from pyspark.sql.functions import lit
from pyspark.sql.types import MapType, StringType, IntegerType, ArrayType, StructType, StructField
import pyspark.sql.functions as F
import random
random_seed = 901
random.seed(random_seed)
spark = SparkSession.builder.appName("Simple test") \
.config("spark.driver.memory", "16g").config("spark.executor.memory", "16g") \
.config("spark.driver.maxResultSize", "0") \
.getOrCreate()
df = spark.createDataFrame(
[
("2009-07-14 08:38:31",11757,"ABC"),
("2009-10-06 10:13:45",34941,"ABC"),
("2009-03-10 09:31:00",40285,"ABC"),
("2009-10-15 08:42:37",41041,"ABC"),
("2009-10-26 10:52:48",49115,"ABC"),
("2009-10-26 13:43:30",49435,"ABC"),
("2009-11-04 14:04:33",57809,"ABC"),
("2009-11-12 00:00:00",66343,"ABC"),
("2009-11-18 12:47:27",72714,"ABC"),
("2009-11-20 10:27:09",74859,"ABC"),
("2009-11-20 13:38:22",75241,"ABC"),
("2009-11-25 14:36:59",80216,"ABC"),
("2009-11-26 15:14:49",81691,"ABC"),
("2009-11-27 14:56:52",83128,"ABC"),
("2009-12-02 08:18:56",86670,"ABC"),
("2009-12-08 13:07:56",93710,"ABC"),
("2009-12-15 12:34:20",101225,"ABC"),
("2009-12-17 08:58:19",103576,"ABC"),
("2009-12-17 10:23:30",103806,"ABC"),
("2010-01-12 12:46:38",118869,"ABC"),
("2010-01-12 13:57:17",119030,"ABC"),
("2010-01-18 11:25:59",124378,"ABC")
],
["revision_date", "revision_id", "cat_name"]
)
# Perform Stratified split. I know that stratified split doesn't make sense for a single `cat_name` type in the toy dataset here but I am preparaing the pipeline for actual dataset.
fractions1 = df.select("cat_name").distinct().withColumn("fraction", F.lit(0.6)).rdd.collectAsMap()
stage_name_train_df = df.stat.sampleBy("cat_name", fractions1, random_seed)
rest_df = df.exceptAll(stage_name_train_df)
fractions2 = {key: 0.5 for key, value in fractions1.items()}
stage_name_valid_df = rest_df.stat.sampleBy("cat_name", fractions2, random_seed)
stage_name_test_df = rest_df.exceptAll(stage_name_valid_df)
stage_name_valid_df.show() # This works fine.
stage_name_valid_df.toPandas() # This works fine as well.
stage_name_valid_df.count() # This gives error.

