org.apache.spark.sql.delta.util.AnalysisHelper.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of wasp-delta-lake_2.12 Show documentation
wasp-delta-lake
The newest version!
/*
 * Copyright (2020) The Delta Lake Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.sql.delta.util

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.analysis.Analyzer
import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.QueryExecutionMetering
import org.apache.spark.sql.delta.DeltaErrors
import org.apache.spark.sql.delta.util.AnalysisHelper.`Amzn-1ReflectionAdapter`.{executeMethodSymbolOpt, executeReflectiveCall}

trait AnalysisHelper {

  import AnalysisHelper._

  protected def tryResolveReferences(
                                      sparkSession: SparkSession
                                    )(expr: Expression, planContainingExpr: LogicalPlan): Expression = {
    val newPlan = FakeLogicalPlan(expr, planContainingExpr.children)
    executionPlan(newPlan)(sparkSession.sessionState.analyzer) match {
      case FakeLogicalPlan(resolvedExpr, _) =>
        // Return even if it did not successfully resolve
        return resolvedExpr
      case _ =>
        // This is unexpected
        throw DeltaErrors.analysisException(s"Could not resolve expression $expr", plan = Option(planContainingExpr))
    }
  }
}

object AnalysisHelper {

  /** LogicalPlan to help resolve the given expression */
  case class FakeLogicalPlan(expr: Expression, children: Seq[LogicalPlan]) extends LogicalPlan {
    override def output: Seq[Attribute] = Nil
  }

  private[util] def executionPlan(newPlan: LogicalPlan)(analyzer: Analyzer): LogicalPlan =
    executeMethodSymbolOpt.map(executeReflectiveCall(newPlan)(analyzer))
      .getOrElse(analyzer.execute(newPlan))

  private[util] object `Amzn-1ReflectionAdapter` {
    import scala.reflect.runtime.{universe => ru}

    def executeMethodSymbolOpt: Option[ru.Symbol] =
      ru
        .typeOf[Analyzer]
        .decl(ru.TermName("execute"))
        .asTerm
        .alternatives
        .find(
          _.asMethod.paramLists.headOption.exists { alternative =>
            alternative.size == 2 &&
            alternative.zip(
              List(
                ru.typeOf[LogicalPlan],
                ru.typeOf[Option[QueryExecutionMetering]]
              )
            ).forall { case (s, typ) => s.typeSignature =:= typ }
          }
        )

    def executeReflectiveCall(
      newPlan: LogicalPlan
    )(analyzer: Analyzer): ru.Symbol => LogicalPlan =
      symbol =>
        ru.runtimeMirror(getClass.getClassLoader)
          .reflect(analyzer)
          .reflectMethod(symbol.asMethod)(newPlan, None)
          .asInstanceOf[LogicalPlan]
  }
}