[BUG] GPU Reader fails due to partition column creating column larger then cudf column size limit

[tgravescs]

Describe the bug
Readers trying to create partition value column hit the cudf column size limit.

 ai.rapids.cudf.CudfException: CUDF failure at: /home/jenkins/agent/workspace/jenkins-spark-rapids-jni_nightly-dev-511-cuda11/thirdparty/cudf/cpp/include/cudf/detail/sizes_to_offsets_iterator.cuh:320: Size of output exceeds the column size limit
	at ai.rapids.cudf.ColumnVector.fromScalar(Native Method)
	at ai.rapids.cudf.ColumnVector.fromScalar(ColumnVector.java:430)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$6(GpuMultiFileReader.scala:264)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$6$adapted(GpuMultiFileReader.scala:263)
	at com.nvidia.spark.rapids.Arm$.withResource(Arm.scala:29)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$5(GpuMultiFileReader.scala:263)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$5$adapted(GpuMultiFileReader.scala:261)
	at scala.collection.IndexedSeqOptimized.foreach(IndexedSeqOptimized.scala:36)
	at scala.collection.IndexedSeqOptimized.foreach$(IndexedSeqOptimized.scala:33)
	at scala.collection.mutable.ArrayOps$ofRef.foreach(ArrayOps.scala:198)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$4(GpuMultiFileReader.scala:261)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$4$adapted(GpuMultiFileReader.scala:260)
	at com.nvidia.spark.rapids.Arm$.withResource(Arm.scala:56)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$3(GpuMultiFileReader.scala:260)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$3$adapted(GpuMultiFileReader.scala:258)
	at scala.collection.TraversableLike$WithFilter.$anonfun$foreach$1(TraversableLike.scala:985)
	at scala.collection.immutable.List.foreach(List.scala:431)
	at scala.collection.TraversableLike$WithFilter.foreach(TraversableLike.scala:984)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$buildPartitionsColumns$1(GpuMultiFileReader.scala:258)
	at com.nvidia.spark.rapids.Arm$.closeOnExcept(Arm.scala:116)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.buildPartitionsColumns(GpuMultiFileReader.scala:257)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.$anonfun$concatAndAddPartitionColsToBatchAndClose$1(GpuMultiFileReader.scala:246)
	at com.nvidia.spark.rapids.Arm$.withResource(Arm.scala:29)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.concatAndAddPartitionColsToBatchAndClose(GpuMultiFileReader.scala:245)
	at com.nvidia.spark.rapids.MultiFileReaderUtils$.addMultiplePartitionValuesAndClose(GpuMultiFileReader.scala:217)
	at com.nvidia.spark.rapids.GpuColumnarBatchWithPartitionValuesIterator.next(GpuColumnarBatchIterator.scala:122)
	at com.nvidia.spark.rapids.GpuColumnarBatchWithPartitionValuesIterator.next(GpuColumnarBatchIterator.scala:101)
	at scala.collection.TraversableOnce$FlattenOps$$anon$2.next(TraversableOnce.scala:522)
	at com.nvidia.spark.rapids.FilePartitionReaderBase.get(GpuMultiFileReader.scala:383)
	at com.nvidia.spark.rapids.FilePartitionReaderBase.get(GpuMultiFileReader.scala:374)
	at com.nvidia.spark.rapids.PartitionIterator.next(dataSourceUtil.scala:39)
	at com.nvidia.spark.rapids.MetricsBatchIterator.next(dataSourceUtil.scala:49)
	at com.nvidia.spark.rapids.MetricsBatchIterator.next(dataSourceUtil.scala:43)
	at com.nvidia.spark.rapids.shims.GpuDataSourceRDD$$anon$1.next(GpuDataSourceRDD.scala:69)
	at org.apache.spark.InterruptibleIterator.next(InterruptibleIterator.scala:40)
	at org.apache.spark.sql.rapids.GpuFileSourceScanExec$$anon$1.next(GpuFileSourceScanExec.scala:483)
	at org.apache.spark.sql.rapids.GpuFileSourceScanExec$$anon$1.next(GpuFileSourceScanExec.scala:472)
	at com.nvidia.spark.rapids.AbstractProjectSplitIterator.next(basicPhysicalOperators.scala:241)
	at com.nvidia.spark.rapids.AbstractProjectSplitIterator.next(basicPhysicalOperators.scala:221)
	at scala.collection.Iterator$$anon$10.next(Iterator.scala:461)
	at scala.collection.Iterator$$anon$10.next(Iterator.scala:461)
	at scala.collection.Iterator$$anon$11.nextCur(Iterator.scala:486)
	at scala.collection.Iterator$$anon$11.hasNext(Iterator.scala:492)
	at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:460)
	at com.nvidia.spark.rapids.GpuMergeAggregateIterator.$anonfun$next$2(aggregate.scala:751)
	at scala.Option.getOrElse(Option.scala:189)
	at com.nvidia.spark.rapids.GpuMergeAggregateIterator.next(aggregate.scala:749)
	at com.nvidia.spark.rapids.GpuMergeAggregateIterator.next(aggregate.scala:711)
	at scala.collection.Iterator$$anon$10.next(Iterator.scala:461)
	at com.nvidia.spark.rapids.DynamicGpuPartialSortAggregateIterator.$anonfun$next$6(aggregate.scala:2035)
	at scala.Option.map(Option.scala:230)
	at com.nvidia.spark.rapids.DynamicGpuPartialSortAggregateIterator.next(aggregate.scala:2035)
	at com.nvidia.spark.rapids.DynamicGpuPartialSortAggregateIterator.next(aggregate.scala:1899)
	at org.apache.spark.sql.rapids.execution.GpuShuffleExchangeExecBase$$anon$1.partNextBatch(GpuShuffleExchangeExecBase.scala:318)
	at org.apache.spark.sql.rapids.execution.GpuShuffleExchangeExecBase$$anon$1.hasNext(GpuShuffleExchangeExecBase.scala:340)
	at org.apache.spark.shuffle.sort.BypassMergeSortShuffleWriter.write(BypassMergeSortShuffleWriter.java:151)
	at org.apache.spark.shuffle.ShuffleWriteProcessor.write(ShuffleWriteProcessor.scala:59)
	at org.apache.spark.scheduler.ShuffleMapTask.$anonfun$runTask$3(ShuffleMapTask.scala:81)
	at com.databricks.spark.util.ExecutorFrameProfiler$.record(ExecutorFrameProfiler.scala:110)
	at org.apache.spark.scheduler.ShuffleMapTask.$anonfun$runTask$1(ShuffleMapTask.scala:81)
	at com.databricks.spark.util.ExecutorFrameProfiler$.record(ExecutorFrameProfiler.scala:110)
	at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:53)
	at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:41)
	at org.apache.spark.scheduler.Task.doRunTask(Task.scala:174)
	at org.apache.spark.scheduler.Task.$anonfun$run$5(Task.scala:142)
	at com.databricks.unity.UCSEphemeralState$Handle.runWith(UCSEphemeralState.scala:41)
	at com.databricks.unity.HandleImpl.runWith(UCSHandle.scala:99)
	at com.databricks.unity.HandleImpl.$anonfun$runWithAndClose$1(UCSHandle.scala:104)
	at scala.util.Using$.resource(Using.scala:269)
	at com.databricks.unity.HandleImpl.runWithAndClose(UCSHandle.scala:103)
	at org.apache.spark.scheduler.Task.$anonfun$run$1(Task.scala:142)
	at com.databricks.spark.util.ExecutorFrameProfiler$.record(ExecutorFrameProfiler.scala:110)
	at org.apache.spark.scheduler.Task.run(Task.scala:97)
	at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$13(Executor.scala:904)
	at org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:1713)
	at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$4(Executor.scala:907)
	at scala.runtime.java8.JFunction0$mcV$sp.apply(JFunction0$mcV$sp.java:23)
	at com.databricks.spark.util.ExecutorFrameProfiler$.record(ExecutorFrameProfiler.scala:110)
	at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:761)
	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)

We should handle this in our readers by checking size and creating multiple batches.

[tgravescs]

Note this happened on Databricks 12.2 but don't think it matters

[revans2]

Once #9085 is in we can probably fix this by adding a split and retry around the code that generates the new column being added in.

[jlowe]

Not sure we need to leverage that, as this is a case of blowing out a scalar to a column. IIRC Spark doesn't support complex types for partition columns, so it should be trivial to properly compute up front how many rows we can afford to produce for each partition column and thus avoid wasted compute, exception catch, and retry with split (especially if retry isn't smart about how it splits).

[revans2]

That would be a better solution to what I was proposing, but we still also want the retry in there just in case we run out of memory because it is not covered by even a simple retry yet.