dfWithFold is not cached, so each fold’s training and validation DataFrames will independently re-evaluate the rand(getSeed) expression and re-scan the input. Persisting dfWithFold once (and unpersisting after building splits) would both reduce work and ensure the fold assignment is identical across all derived DataFrames even if the optimizer rewrites the plan.

In non-barrier mode this uses df.coalesce(numTasks), but coalesce is a no-op when numTasks is greater than the current partition count. Combined with the updated determineNumTasks, this can leave the DataFrame with fewer partitions than numTasks, while the training logic still assumes numTasks workers. Consider guarding against partition increases here (or using repartition(numTasks) when increasing partitions) so the DataFrame partitioning matches the worker count used for coordination.

This issue also appears on line 444 of the same file.

numPartitions is derived via max(spark_partition_id()), which can return null on an empty DataFrame; head().getInt(0) will then throw. Also, using max(partitionId) + 1 can undercount when trailing partitions are empty (partition IDs are not guaranteed to be contiguous in the data). Please handle the empty-input case safely and avoid assuming maxPid + 1 equals the real partition count.

This issue also appears on line 100 of the same file.

The round-robin index calculation looks off by one: row_number() starts at 1, so row_number() % numPartitions yields values in 1..numPartitions-1 with 0 only every numPartitions rows, which skews distribution. Additionally, repartitionByRange(numPartitions, col("_rr_idx")) does not guarantee that each _rr_idx value maps to a single partition (range boundaries are sampled), so it may not achieve the intended per-partition label spread. Consider using a zero-based modulo (e.g., pmod(row_number()-1, numPartitions)) and repartition(numPartitions, col("_rr_idx")) if you need deterministic partition assignment by _rr_idx.

Switching from conditional coalesce/RDD-based repartitioning to unconditional dataset.repartition(getN) means reducing partitions will now always incur a full shuffle (previously it could use coalesce without shuffle when decreasing partitions). If this transformer is expected to be used for downsizing partitions, consider adding a parameter to choose coalesce vs repartition, or reintroduce a non-shuffle path when getN is less than the current partition count (without relying on RDD APIs).

monotonically_increasing_id() does not produce dense, 0-based consecutive indices (it encodes partition ID and can create very large gaps). Downstream MatrixOps.size computes matrix dimensions via max(i)+1, so this change can massively inflate inferred sizes and lead to incorrect behavior or OOMs. Please generate a dense index instead (e.g., row_number() over an explicit Window.orderBy(...) and subtract 1) so indices are contiguous and bounded by the distinct count.

SparkSession.active will throw if there is no active Spark session on the current thread. DefaultCacheDir is a lazy val and can be initialized outside of Spark execution (or before SparkSession.setActiveSession), which would make ONNXHub fail unexpectedly. Consider using SparkSession.builder().getOrCreate() or SparkSession.getActiveSession.getOrElse(SparkSession.builder().getOrCreate()) to obtain a session more safely here.

This issue also appears on line 196 of the same file.

See below for a potential fix:

        val spark = SparkSession.getActiveSession.getOrElse(SparkSession.builder().getOrCreate())
        val home = new Path("placeholder")
          .getFileSystem(spark.sparkContext.hadoopConfiguration)

Converting getLen (Long) to Int can overflow for model files >2GB, leading to negative/incorrect array sizes and corrupted reads. Consider handling lengths > Int.MaxValue explicitly (e.g., fail with a clear error) or reading the stream without pre-allocating an Int-sized byte array.

In non-barrier mode, dfSubset.coalesce(numExecutorTasks) will not increase the number of partitions, but this method still returns numExecutorTasks as numTasks. That can make downstream coordination (e.g., augmentVowpalWabbitArguments, --node numbering) think there are more tasks than Spark will actually run, which can lead to hangs/timeouts or incorrect distributed setup. Consider ensuring the returned task count matches the actual partition count (e.g., by using repartition(numExecutorTasks) when increasing partitions, or by deriving the current partition count from the DataFrame plan and capping numExecutorTasks).