Github user sabanas commented on a diff in the pull request:
https://github.com/apache/spark/pull/20915#discussion_r178090546
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala
---
@@ -50,6 +51,85 @@ import org.apache.spark.sql.execution.SparkPlan
* and add it. Proceed to the next file.
*/
object FileSourceStrategy extends Strategy with Logging {
+
+ // should prune buckets iff num buckets is greater than 1 and there is
only one bucket column
+ private def shouldPruneBuckets(bucketSpec: Option[BucketSpec]): Boolean
= {
+ bucketSpec match {
+ case Some(spec) => spec.bucketColumnNames.length == 1 &&
spec.numBuckets > 1
+ case None => false
+ }
+ }
+
+ private def getExpressionBuckets(expr: Expression,
+ bucketColumnName: String,
+ numBuckets: Int): BitSet = {
+
+ def getMatchedBucketBitSet(attr: Attribute, v: Any): BitSet = {
+ val matchedBuckets = new BitSet(numBuckets)
+ matchedBuckets.set(BucketingUtils.getBucketIdFromValue(attr,
numBuckets, v))
+ matchedBuckets
+ }
+
+ expr match {
+ case expressions.EqualTo(a: Attribute, Literal(v, _)) if a.name ==
bucketColumnName =>
+ getMatchedBucketBitSet(a, v)
+ case expressions.EqualTo(Literal(v, _), a: Attribute) if a.name ==
bucketColumnName =>
+ getMatchedBucketBitSet(a, v)
+ case expressions.EqualNullSafe(a: Attribute, Literal(v, _)) if
a.name == bucketColumnName =>
+ getMatchedBucketBitSet(a, v)
+ case expressions.EqualNullSafe(Literal(v, _), a: Attribute) if
a.name == bucketColumnName =>
+ getMatchedBucketBitSet(a, v)
+ case expressions.In(a: Attribute, list)
+ if list.forall(_.isInstanceOf[Literal]) && a.name ==
bucketColumnName =>
+ val valuesSet = list.map(e => e.eval(EmptyRow))
+ valuesSet
+ .map(v => getMatchedBucketBitSet(a, v))
+ .fold(new BitSet(numBuckets))(_ | _)
+ case expressions.IsNull(a: Attribute) if a.name == bucketColumnName
=>
+ getMatchedBucketBitSet(a, null)
+ case expressions.And(left, right) =>
+ getExpressionBuckets(left, bucketColumnName, numBuckets) |
--- End diff --
in an AND condition i'd like to get the union of both RHS and LHS BitSets.
for example, if table is bucketed on column j:
j == 0 && i == 0 -> will return a bit set matching j's bucket.
i == 0 && k == 0 -> will return union of 2 empty bit sets, meaning an empty
bit set.
j == 0 && j == 1 (which is effectively a false condition)-> will return bit
set matching both j values. i haven't implemented an optimization for this case
(can be done with & condition in case both RHS & LHS are not empty).
can you think of an input that will not work here?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
