DataSourceStrategy

Author: dbtsai

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala

@@ -442,53 +442,65 @@ object DataSourceStrategy {
    * @return a `Some[Filter]` if the input [[Expression]] is convertible, otherwise a `None`.
    */
   protected[sql] def translateFilter(predicate: Expression): Option[Filter] = {
+    // Recursively try to find an attribute name from the top level that can be pushed down.
+    def attrName(e: Expression): Option[String] = e match {
+      // In Spark and many data sources such as parquet, dots are used as a column path delimiter;
+      // thus, we don't translate such expressions.
+      case a: Attribute if !a.name.contains(".") =>
+        Some(a.name)
+      case s: GetStructField if !s.childSchema(s.ordinal).name.contains(".") =>
+        attrName(s.child).map(_ + s".${s.childSchema(s.ordinal).name}")
+      case _ =>
+        None
+    }
+
     predicate match {
-      case expressions.EqualTo(a: Attribute, Literal(v, t)) =>
-        Some(sources.EqualTo(a.name, convertToScala(v, t)))
-      case expressions.EqualTo(Literal(v, t), a: Attribute) =>
-        Some(sources.EqualTo(a.name, convertToScala(v, t)))
-
-      case expressions.EqualNullSafe(a: Attribute, Literal(v, t)) =>
-        Some(sources.EqualNullSafe(a.name, convertToScala(v, t)))
-      case expressions.EqualNullSafe(Literal(v, t), a: Attribute) =>
-        Some(sources.EqualNullSafe(a.name, convertToScala(v, t)))
-
-      case expressions.GreaterThan(a: Attribute, Literal(v, t)) =>
-        Some(sources.GreaterThan(a.name, convertToScala(v, t)))
-      case expressions.GreaterThan(Literal(v, t), a: Attribute) =>
-        Some(sources.LessThan(a.name, convertToScala(v, t)))
-
-      case expressions.LessThan(a: Attribute, Literal(v, t)) =>
-        Some(sources.LessThan(a.name, convertToScala(v, t)))
-      case expressions.LessThan(Literal(v, t), a: Attribute) =>
-        Some(sources.GreaterThan(a.name, convertToScala(v, t)))
-
-      case expressions.GreaterThanOrEqual(a: Attribute, Literal(v, t)) =>
-        Some(sources.GreaterThanOrEqual(a.name, convertToScala(v, t)))
-      case expressions.GreaterThanOrEqual(Literal(v, t), a: Attribute) =>
-        Some(sources.LessThanOrEqual(a.name, convertToScala(v, t)))
-
-      case expressions.LessThanOrEqual(a: Attribute, Literal(v, t)) =>
-        Some(sources.LessThanOrEqual(a.name, convertToScala(v, t)))
-      case expressions.LessThanOrEqual(Literal(v, t), a: Attribute) =>
-        Some(sources.GreaterThanOrEqual(a.name, convertToScala(v, t)))
-
-      case expressions.InSet(a: Attribute, set) =>
-        val toScala = CatalystTypeConverters.createToScalaConverter(a.dataType)
-        Some(sources.In(a.name, set.toArray.map(toScala)))
+      case expressions.EqualTo(e: Expression, Literal(v, t)) =>
+        attrName(e).map(name => sources.EqualTo(name, convertToScala(v, t)))
+      case expressions.EqualTo(Literal(v, t), e: Expression) =>
+        attrName(e).map(name => sources.EqualTo(name, convertToScala(v, t)))
+
+      case expressions.EqualNullSafe(e: Expression, Literal(v, t)) =>
+        attrName(e).map(name => sources.EqualNullSafe(name, convertToScala(v, t)))
+      case expressions.EqualNullSafe(Literal(v, t), e: Expression) =>
+        attrName(e).map(name => sources.EqualNullSafe(name, convertToScala(v, t)))
+
+      case expressions.GreaterThan(e: Expression, Literal(v, t)) =>
+        attrName(e).map(name => sources.GreaterThan(name, convertToScala(v, t)))
+      case expressions.GreaterThan(Literal(v, t), e: Expression) =>
+        attrName(e).map(name => sources.LessThan(name, convertToScala(v, t)))
+
+      case expressions.LessThan(e: Expression, Literal(v, t)) =>
+        attrName(e).map(name => sources.LessThan(name, convertToScala(v, t)))
+      case expressions.LessThan(Literal(v, t), e: Expression) =>
+        attrName(e).map(name => sources.GreaterThan(name, convertToScala(v, t)))
+
+      case expressions.GreaterThanOrEqual(e: Expression, Literal(v, t)) =>
+        attrName(e).map(name => sources.GreaterThanOrEqual(name, convertToScala(v, t)))
+      case expressions.GreaterThanOrEqual(Literal(v, t), e: Expression) =>
+        attrName(e).map(name => sources.LessThanOrEqual(name, convertToScala(v, t)))
+
+      case expressions.LessThanOrEqual(e: Expression, Literal(v, t)) =>
+        attrName(e).map(name => sources.LessThanOrEqual(name, convertToScala(v, t)))
+      case expressions.LessThanOrEqual(Literal(v, t), e: Expression) =>
+        attrName(e).map(name => sources.GreaterThanOrEqual(name, convertToScala(v, t)))
+
+      case expressions.InSet(e: Expression, set) =>
+        val toScala = CatalystTypeConverters.createToScalaConverter(e.dataType)
+        attrName(e).map(name => sources.In(name, set.toArray.map(toScala)))
 
       // Because we only convert In to InSet in Optimizer when there are more than certain
       // items. So it is possible we still get an In expression here that needs to be pushed
       // down.
-      case expressions.In(a: Attribute, list) if list.forall(_.isInstanceOf[Literal]) =>
-        val hSet = list.map(_.eval(EmptyRow))
-        val toScala = CatalystTypeConverters.createToScalaConverter(a.dataType)
-        Some(sources.In(a.name, hSet.toArray.map(toScala)))
+      case expressions.In(e: Expression, list) if list.forall(_.isInstanceOf[Literal]) =>
+        val hSet = list.map(e => e.eval(EmptyRow))
+        val toScala = CatalystTypeConverters.createToScalaConverter(e.dataType)
+        attrName(e).map(name => sources.In(name, hSet.toArray.map(toScala)))
 
-      case expressions.IsNull(a: Attribute) =>
-        Some(sources.IsNull(a.name))
-      case expressions.IsNotNull(a: Attribute) =>
-        Some(sources.IsNotNull(a.name))
+      case expressions.IsNull(e: Expression) =>
+        attrName(e).map(name => sources.IsNull(name))
+      case expressions.IsNotNull(e: Expression) =>
+        attrName(e).map(name => sources.IsNotNull(name))
 
       case expressions.And(left, right) =>
         // See SPARK-12218 for detailed discussion
@@ -514,14 +526,14 @@ object DataSourceStrategy {
       case expressions.Not(child) =>
         translateFilter(child).map(sources.Not)
 
-      case expressions.StartsWith(a: Attribute, Literal(v: UTF8String, StringType)) =>
-        Some(sources.StringStartsWith(a.name, v.toString))
+      case expressions.StartsWith(e: Expression, Literal(v: UTF8String, StringType)) =>
+        attrName(e).map(name => sources.StringStartsWith(name, v.toString))
 
-      case expressions.EndsWith(a: Attribute, Literal(v: UTF8String, StringType)) =>
-        Some(sources.StringEndsWith(a.name, v.toString))
+      case expressions.EndsWith(e: Expression, Literal(v: UTF8String, StringType)) =>
+        attrName(e).map(name => sources.StringEndsWith(name, v.toString))
 
-      case expressions.Contains(a: Attribute, Literal(v: UTF8String, StringType)) =>
-        Some(sources.StringContains(a.name, v.toString))
+      case expressions.Contains(e: Expression, Literal(v: UTF8String, StringType)) =>
+        attrName(e).map(name => sources.StringContains(name, v.toString))
 
       case expressions.Literal(true, BooleanType) =>
         Some(sources.AlwaysTrue)