[SPARK-17170] [SQL] InMemoryTableScanExec driver-side partition pruning #14733
Conversation
| if (validPartitions.isEmpty) { | ||
| new EmptyRDD[CachedBatch](sparkContext) | ||
| } else { | ||
| new PartitionPruningRDD[CachedBatch](relation.cachedColumnBuffers, |
There was a problem hiding this comment.
There's PartitionPruningRDD.create which will make this slightly cleaner and if you skip logging it's just
PartitionPruningRDD.create(relation.cachedColumnBuffers, validPartitions.contains).
There was a problem hiding this comment.
Cool yeah used this (and removed the EmptyRDD constructor as well). I'd prefer to keep the logging in the function though.
|
@rxin @davies @dongjoon-hyun mind taking a look? |
| if (validPartitions.contains(index)) { | ||
| true | ||
| } else { | ||
| logInfo(s"Skipping partition $index because all cached batches will be pruned") |
There was a problem hiding this comment.
Current executor-side pruning logging is done at INFO. I have no strong opinion either way, but this can get noisy with many partitions getting pruned.
|
Jenkins, this is ok to test. |
| @@ -106,7 +106,7 @@ case class InMemoryRelation( | |||
|
|
|||
| private def buildBuffers(): Unit = { | |||
| val output = child.output | |||
| val cached = child.execute().mapPartitionsInternal { rowIterator => | |||
| val cached = child.execute().mapPartitionsWithIndex { (i, rowIterator) => | |||
There was a problem hiding this comment.
Here, I'd prefer to use TaskContext.get.getPartitionId to get the partition id. The problem with this change is that it's re-introducing closure cleaning overhead, which mapPartitionsInternal avoids.
|
This looks pretty cool! I'll come back later tonight / tomorrow to test this out and do a more detailed review pass. |
|
Test build #64611 has finished for PR 14733 at commit
|
|
Test build #64654 has finished for PR 14733 at commit
|
|
Test build #64659 has finished for PR 14733 at commit
|
|
This looks like a legitimate test failure: |
| @@ -142,13 +171,16 @@ case class InMemoryTableScanExec( | |||
| val cachedBatchesToScan = | |||
| if (inMemoryPartitionPruningEnabled) { | |||
| cachedBatchIterator.filter { cachedBatch => | |||
| val partitionFilter = newPredicate( | |||
There was a problem hiding this comment.
Why did you choose to nest this more deeply inside of the filter rather than leaving it where it was in mapPartitionsInternal? By moving it here, we'll wind up calling newPredicate once per batch rather than once per partition, thereby harming performance.
There was a problem hiding this comment.
It was to avoid the call if inMemoryPartitionPruning wasn't enabled. This reasoning is kind of dumb though given that it is the default and if disabled you will pay extra cost elsewhere. I'll move it back.
|
I believe that the test failure is happening because I think that this is breaking One possible fix would be to have It might be possible to work around this via a custom, SQL-specific version of ZippedPartitionsRDD which understands how to deal with missing partitions and zips according to partition ids (dealing with gaps), but this will be tricky to get right for outer joins (where you still need to produce output for partitions which match to pruned ones). |
|
Closing stale PR |
What changes were proposed in this pull request?
After caching data, we have statistics that enable us to eagerly prune entire partitions before launching a query. This modifies the InMemoryTableScanExec to prune partitions before launching the tasks.
How was this patch tested?
Existing test suite with slight modification to scan over the data once as setup.