Update partitioning logic in ShuffledBatchRDD

[andygrove]

This PR updates the logic in ShuffledBatchRDD to reflect recent changes in Spark's ShuffledRowRDD related to AQE support.

Note that there are oustanding questions about how these changes affect UCX, so UCX is disabled for now when AQE is enabled.

[andygrove]

@abellina Could you take a look when you get a chance? This is one of the changes I had to make when working on the AQE POC.

[andygrove]

build

[kuhushukla]

I'll try and chime in this weekend on this change.

[abellina]

To me at high level this looks fine. It is echoing the changes made in the row-based ShuffledRowRDD. The question I have is how it works with the shuffle plugin. Given AQE, with the shuffle plugin, the expectation is that this will likely fail, so we may want to turn off the shuffle plugin in that case (it could be a separate bug). I'll test with your patch (AQE) and the shuffle plugin also.

[abellina]

@andygrove I haven't had time to test this yet, but I am fairly sure it will fail if the shuffle plugin is enabled, I'll try this today with some help from you likely (not sure what I can run with it enabled).

Here's what I am thinking: the writes are going to be cached and stamped as rapids blocks, and then the reads are going to ignore this all together (the getReaderForRange falls back as if it were a CPU shuffle). We do need to implement getReaderForRange for the plugin, but one way to get around this is to revert back to the legacy shuffle if AQE is on for now (https://github.com/NVIDIA/spark-rapids/blob/branch-0.2/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RapidsShuffleInternalManager.scala#L194).

[abellina]

@andygrove here's the issue I mentioned on the read side: #362. I think it's a separate issue, that should come with tests also (somehow)

[andygrove]

build

[revans2]

There is a stringDefault value in SQLConf.ADAPTIVE_EXECUTION_ENABLED which should let us stay up to date with the default if it changes. The main problem with that is that we are compiling this against one specific version of spark so if the default changes from one version to another this will not stay up to date.

[andygrove]

This is something that we could potentially do in the shim layer though so that it does compile against each version we support. I'll look into doing that.

[abellina]

@andygrove we should wait on this PR then? Or you are thinking a different PR?

[andygrove]

I'll add it to this PR.

[andygrove]

Actually, this is a boolean config so it has to have a value and it doesn't make sense to check for a default value. I'll remove the default part instead.

[andygrove]

This is updated now.

[andygrove]

build

[andygrove]

build

[andygrove]

There was a build failure with the 3.1.0 shim:

13:44:04  [ERROR] case class GpuBroadcastHashJoinExec(
13:44:04  [ERROR]            ^
13:44:04  [ERROR] /ansible-managed/jenkins-slave/slave4/workspace/spark/rapids_premerge-github/shims/spark310/src/main/scala/com/nvidia/spark/rapids/shims/spark310/GpuShuffledHashJoinExec.scala:77: class GpuShuffledHashJoinExec needs to be abstract, since:
13:44:04  it has 2 unimplemented members.
13:44:04  /** As seen from class GpuShuffledHashJoinExec, the missing signatures are as follows.
13:44:04   *  For convenience, these are usable as stub implementations.
13:44:04   */
13:44:04    // Members declared in org.apache.spark.sql.execution.CodegenSupport
13:44:04    def inputRDDs(): Seq[org.apache.spark.rdd.RDD[org.apache.spark.sql.catalyst.InternalRow]] = ???
13:44:04    
13:44:04    // Members declared in org.apache.spark.sql.execution.joins.HashJoin
13:44:04    protected def prepareRelation(ctx: org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext): (String, Boolean) = ???

Seems to be related to this recent Spark commit:

apache/spark@ae82768

[andygrove]

build