Spark DataSource 数据源

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
/**
   * Loads input in as a `DataFrame`, for data sources that support multiple paths.
   * Only works if the source is a HadoopFsRelationProvider.
   *
   * @since 1.6.0
   */
@scala.annotation.varargs
def load(paths: String*): DataFrame = {
  if (source.toLowerCase(Locale.ROOT) == DDLUtils.HIVE_PROVIDER) {
    throw new AnalysisException("Hive data source can only be used with tables, you can not " +
                                "read files of Hive data source directly.")
 }

  val cls = DataSource.lookupDataSource(source, sparkSession.sessionState.conf)
  // DataSourceV2
  if (classOf[DataSourceV2].isAssignableFrom(cls)) {
    val ds = cls.newInstance().asInstanceOf[DataSourceV2]
    if (ds.isInstanceOf[ReadSupport]) {
      val sessionOptions = DataSourceV2Utils.extractSessionConfigs(
        ds = ds, conf = sparkSession.sessionState.conf)
      val pathsOption = {
        val objectMapper = new ObjectMapper()
        DataSourceOptions.PATHS_KEY -> objectMapper.writeValueAsString(paths.toArray)
     }
      Dataset.ofRows(sparkSession, DataSourceV2Relation.create(
        ds, sessionOptions ++ extraOptions.toMap + pathsOption,
        userSpecifiedSchema = userSpecifiedSchema))
   } else {
      loadV1Source(paths: _*)
   }
 } else {
    loadV1Source(paths: _*)
 }
}

private def loadV1Source(paths: String*) = {
  // Code path for data source v1.
  // baseRelationToDataFrame
  sparkSession.baseRelationToDataFrame(
    DataSource.apply(
      sparkSession,
      paths = paths,
      userSpecifiedSchema = userSpecifiedSchema,
      className = source,
      options = extraOptions.toMap)
    // resolveRelation
   .resolveRelation())
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
private[sql] case class JDBCRelation(
    override val schema: StructType,
    parts: Array[Partition],
    jdbcOptions: JDBCOptions)(@transient val sparkSession: SparkSession)
  extends BaseRelation
  with PrunedFilteredScan
  with InsertableRelation {
    
    // ...
    
  override def buildScan(requiredColumns: Array[String], filters: Array[Filter]): RDD[Row] = {
    // Rely on a type erasure hack to pass RDD[InternalRow] back as RDD[Row]
    JDBCRDD.scanTable(
      sparkSession.sparkContext,
      schema,
      requiredColumns,
      filters,
      parts,
      jdbcOptions).asInstanceOf[RDD[Row]]
 }

  override def insert(data: DataFrame, overwrite: Boolean): Unit = {
    data.write
     .mode(if (overwrite) SaveMode.Overwrite else SaveMode.Append)
     .jdbc(jdbcOptions.url, jdbcOptions.tableOrQuery, jdbcOptions.asProperties)
 }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
// object JDBCRDD
def scanTable(
  sc: SparkContext,
  schema: StructType,
  requiredColumns: Array[String],
  filters: Array[Filter],
  parts: Array[Partition],
  options: JDBCOptions): RDD[InternalRow] = {
  val url = options.url
  val dialect = JdbcDialects.get(url)
  val quotedColumns = requiredColumns.map(colName => dialect.quoteIdentifier(colName))
  // class JDBCRDD
  new JDBCRDD(
    sc,
    JdbcUtils.createConnectionFactory(options),
    pruneSchema(schema, requiredColumns),
    quotedColumns,
    filters,
    parts,
    url,
    options)
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class JdbcRelationProvider extends CreatableRelationProvider
  with RelationProvider with DataSourceRegister {
  
  override def shortName(): String = "jdbc"  
  
  // RelationProvider
  override def createRelation(
      sqlContext: SQLContext,
      parameters: Map[String, String]): BaseRelation = {
    val jdbcOptions = new JDBCOptions(parameters)
    val resolver = sqlContext.conf.resolver
    val timeZoneId = sqlContext.conf.sessionLocalTimeZone
    // schema
    val schema = JDBCRelation.getSchema(resolver, jdbcOptions)
    // 分区
    val parts = JDBCRelation.columnPartition(schema, resolver, timeZoneId, jdbcOptions)
    JDBCRelation(schema, parts, jdbcOptions)(sqlContext.sparkSession)
 }
    
  // CreatableRelationProvider
  override def createRelation(
      sqlContext: SQLContext,
      mode: SaveMode,
      parameters: Map[String, String],
      df: DataFrame): BaseRelation = {
    val options = new JdbcOptionsInWrite(parameters)
    val isCaseSensitive = sqlContext.conf.caseSensitiveAnalysis
    val conn = JdbcUtils.createConnectionFactory(options)()
    // 判断表是否存在，如果存在判断写入模式（Overwrite、Append、ErrorIfExists、Ignore）作不同处理
    ....

    createRelation(sqlContext, parameters)
 }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
private[jdbc] class JDBCRDD(
    sc: SparkContext,
    getConnection: () => Connection,
    schema: StructType,
    columns: Array[String],
    filters: Array[Filter],
    partitions: Array[Partition],
    url: String,
    options: JDBCOptions)
  extends RDD[InternalRow](sc, Nil) {
    
  override def compute(thePart: Partition, context: TaskContext): Iterator[InternalRow] = {
      // ...
      
    // 根据filters和partition构造where条件
    val myWhereClause = getWhereClause(part)
    
    // 生成对应sql并执行
    val sqlText = s"SELECT $columnList FROM ${options.table} $myWhereClause"
    stmt = conn.prepareStatement(sqlText,
                                 ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)
    stmt.setFetchSize(options.fetchSize)
    rs = stmt.executeQuery()
    val rowsIterator = JdbcUtils.resultSetToSparkInternalRows(rs, schema, inputMetrics)

    CompletionIterator[InternalRow, Iterator[InternalRow]](
      new InterruptibleIterator(context, rowsIterator), close())
 }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
// DataSourceV2
public interface DataSourceV2 {}

// ReadSupport
public interface ReadSupport extends DataSourceV2 {
  DataSourceReader createReader(DataSourceOptions options);
}

// DataSourceReader
public interface DataSourceReader {
  StructType readSchema();
  List<DataReaderFactory<Row>> createDataReaderFactories();
}

// DataReaderFactory
public interface DataReaderFactory<T> extends Serializable {
  default String[] preferredLocations() {
    return new String[0];
 }
  DataReader<T> createDataReader();
}

// DataReader
public interface DataReader<T> extends Closeable {
  boolean next() throws IOException;
  T get();
}

1
2
3
4
5
6
7
8
9
10
11
12
// SupportsPushDownRequiredColumns
// 在 PushDownOperatorsToDataSource 中会调用 pruneColumns() 方法
public interface SupportsPushDownRequiredColumns extends DataSourceReader {
  void pruneColumns(StructType requiredSchema);
}

// SupportsPushDownFilters
// 在 PushDownOperatorsToDataSource 中会调用 pushFilters() 方法
public interface SupportsPushDownFilters extends DataSourceReader {
  Filter[] pushFilters(Filter[] filters);
  Filter[] pushedFilters();
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def load(paths: String*): DataFrame = {
  if (source.toLowerCase(Locale.ROOT) == DDLUtils.HIVE_PROVIDER) {
    throw new AnalysisException("Hive data source can only be used with tables, you can not " +
                                "read files of Hive data source directly.")
 }

  val cls = DataSource.lookupDataSource(source, sparkSession.sessionState.conf)
  // DataSourceV2
  if (classOf[DataSourceV2].isAssignableFrom(cls)) {
    val ds = cls.newInstance()
    val options = new DataSourceOptions((extraOptions ++
                                         DataSourceV2Utils.extractSessionConfigs(
                                           ds = ds.asInstanceOf[DataSourceV2],
                                           conf = sparkSession.sessionState.conf)).asJava)

    // DataSourceReader
    val reader = (ds, userSpecifiedSchema) match {
      case (ds: ReadSupportWithSchema, Some(schema)) =>
      ds.createReader(schema, options)

      case (ds: ReadSupport, None) =>
      ds.createReader(options)

      case (ds: ReadSupportWithSchema, None) =>
      throw new AnalysisException(s"A schema needs to be specified when using $ds.")

      case (ds: ReadSupport, Some(schema)) =>
      val reader = ds.createReader(options)
      if (reader.readSchema() != schema) {
        throw new AnalysisException(s"$ds does not allow user-specified schemas.")
     }
      reader

      case _ => null // fall back to v1
   }

    if (reader == null) {
      loadV1Source(paths: _*)
   } else {
      // reader -> DataSourceV2Relation -> Dataset
      Dataset.ofRows(sparkSession, DataSourceV2Relation(reader))
   }
 } else {
    loadV1Source(paths: _*)
 }
}

1
2
3
// DataSourceV2Relation apply():
//   readSchema()
new DataSourceV2Relation(reader.readSchema().toAttributes, reader)

1
2
3
4
5
6
7
8
9
10
11
object DataSourceV2Strategy extends Strategy {
  override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
    case DataSourceV2Relation(output, reader) =>
      DataSourceV2ScanExec(output, reader) :: Nil

    case WriteToDataSourceV2(writer, query) =>
      WriteToDataSourceV2Exec(writer, planLater(query)) :: Nil

    case _ => Nil
 }
}

1
2
3
4
5
6
7
8
9
10
11
// createDataReaderFactories() -> DataReaderFactory
private lazy val readerFactories: java.util.List[DataReaderFactory[UnsafeRow]] = reader match {
  case r: SupportsScanUnsafeRow => r.createUnsafeRowReaderFactories()
  case _ =>
  reader.createDataReaderFactories().asScala.map {
    new RowToUnsafeRowDataReaderFactory(_, reader.readSchema()): DataReaderFactory[UnsafeRow]
 }.asJava
}

// inputRDD
new DataSourceRDD(sparkContext, readerFactories).asInstanceOf[RDD[InternalRow]]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
override def compute(split: Partition, context: TaskContext): Iterator[T] = {
  val reader = split.asInstanceOf[DataSourceRDDPartition[T]].readerFactory.createDataReader()
  context.addTaskCompletionListener(_ => reader.close())
  val iter = new Iterator[T] {
    private[this] var valuePrepared = false

    override def hasNext: Boolean = {
      if (!valuePrepared) {
        // 调用 DataReader next()
        valuePrepared = reader.next()
     }
      valuePrepared
   }

    override def next(): T = {
      if (!hasNext) {
        throw new java.util.NoSuchElementException("End of stream")
     }
      valuePrepared = false
      // 调用 DataReader get()
      reader.get()
   }
 }
  new InterruptibleIterator(context, iter)
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
// DataSourceV2
public interface DataSourceV2 {}

// ReadSupport
public interface ReadSupport extends DataSourceV2 {
  // 新增带schema的方式，需自定义重写
  default DataSourceReader createReader(StructType schema, DataSourceOptions options) {
    String name;
    if (this instanceof DataSourceRegister) {
      name = ((DataSourceRegister) this).shortName();
    } else {
      name = this.getClass().getName();
    }
    // 默认抛出异常
    throw new UnsupportedOperationException(name + " does not support user specified schema");
  }
  DataSourceReader createReader(DataSourceOptions options);
}

// DataSourceReader
public interface DataSourceReader {
  StructType readSchema();
  List<InputPartition<InternalRow>> planInputPartitions();
}

// InputPartition
public interface InputPartition<T> extends Serializable {
  default String[] preferredLocations() {
    return new String[0];
  }
  InputPartitionReader<T> createPartitionReader();
}

// InputPartitionReader
public interface InputPartitionReader<T> extends Closeable {
  boolean next() throws IOException;
  T get();
}

1
2
3
4
5
6
7
8
9
10
def enableHiveSupport(): Builder = synchronized {
  if (hiveClassesArePresent) {
    // 配置 CATALOG_IMPLEMENTATION
    config(CATALOG_IMPLEMENTATION.key, "hive")
 } else {
    throw new IllegalArgumentException(
      "Unable to instantiate SparkSession with Hive support because " +
      "Hive classes are not found.")
 }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
private val HIVE_SESSION_STATE_BUILDER_CLASS_NAME =
"org.apache.spark.sql.hive.HiveSessionStateBuilder"

private def sessionStateClassName(conf: SparkConf): String = {
  // 获取 CATALOG_IMPLEMENTATION
  conf.get(CATALOG_IMPLEMENTATION) match {
    case "hive" => HIVE_SESSION_STATE_BUILDER_CLASS_NAME
    case "in-memory" => classOf[SessionStateBuilder].getCanonicalName
 }
}

@Unstable
@transient
lazy val sessionState: SessionState = {
  parentSessionState
 .map(_.clone(this))
 .getOrElse {
    // 初始化 sessionState
    val state = SparkSession.instantiateSessionState(
      SparkSession.sessionStateClassName(sharedState.conf),
      self)
    state
 }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
private[sql] class HiveSessionCatalog(
    externalCatalogBuilder: () => ExternalCatalog,
    globalTempViewManagerBuilder: () => GlobalTempViewManager,
    val metastoreCatalog: HiveMetastoreCatalog,
    functionRegistry: FunctionRegistry,
    tableFunctionRegistry: TableFunctionRegistry,
    hadoopConf: Configuration,
    parser: ParserInterface,
    functionResourceLoader: FunctionResourceLoader)
  extends SessionCatalog(
      externalCatalogBuilder,
      globalTempViewManagerBuilder,
      functionRegistry,
      tableFunctionRegistry,
      hadoopConf,
      parser,
      functionResourceLoader) {
  // ...
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
/**
 * An internal catalog that is used by a Spark Session. This internal catalog serves as a
 * proxy to the underlying metastore (e.g. Hive Metastore) and it also manages temporary
 * views and functions of the Spark Session that it belongs to.
 *
 * This class must be thread-safe.
 */
class SessionCatalog(
    // ExternalCatalog: [HiveExternalCatalog, InMemoryCatalog]
    externalCatalogBuilder: () => ExternalCatalog,
    globalTempViewManagerBuilder: () => GlobalTempViewManager,
    functionRegistry: FunctionRegistry,
    tableFunctionRegistry: TableFunctionRegistry,
    hadoopConf: Configuration,
    parser: ParserInterface,
    functionResourceLoader: FunctionResourceLoader,
    cacheSize: Int = SQLConf.get.tableRelationCacheSize,
    cacheTTL: Long = SQLConf.get.metadataCacheTTL) extends SQLConfHelper with Logging {
  
  lazy val externalCatalog = externalCatalogBuilder()
  
  def listDatabases(): Seq[String] = {
    // 实际使用 ExternalCatalog
    externalCatalog.listDatabases()
 }
  // ...
}

1
2
3
4
5
6
7
8
9
10
11
12
13
private[spark] class HiveExternalCatalog(conf: SparkConf, hadoopConf: Configuration)
  extends ExternalCatalog with Logging{

  lazy val client: HiveClient = {
    HiveUtils.newClientForMetadata(conf, hadoopConf)
 }

  override def listDatabases(): Seq[String] = withClient {
    // 实际使用 HiveClient
    client.listDatabases("*")
 }
  // ...
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def lookupRelation(name: TableIdentifier): LogicalPlan = {
  synchronized {
    val db = formatDatabaseName(name.database.getOrElse(currentDb))
    val table = formatTableName(name.table)
    // db是globalTempView的默认数据库，则为视图类型
    if (db == globalTempViewManager.database) {
      globalTempViewManager.get(table).map { viewDef =>
        // 视图，TempView
        SubqueryAlias(table, db, getTempViewPlan(viewDef))
     }.getOrElse(throw new NoSuchTableException(db, table))
   } else if (name.database.isDefined || !tempViews.contains(table)) {
      // ExternalCatalog查询表元数据
      val metadata = externalCatalog.getTable(db, table)
      // 生成 Relation
      getRelation(metadata)
   } else {
      // 视图，TempView
      SubqueryAlias(table, getTempViewPlan(tempViews(table)))
   }
 }
}

def getRelation(
  metadata: CatalogTable,
  options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()): LogicalPlan = {
  val name = metadata.identifier
  val db = formatDatabaseName(name.database.getOrElse(currentDb))
  val table = formatTableName(name.table)
  val multiParts = Seq(CatalogManager.SESSION_CATALOG_NAME, db, table)

  // 视图，生成 VIEW
  if (metadata.tableType == CatalogTableType.VIEW) {
    // The relation is a view, so we wrap the relation by:
    // 1. Add a [[View]] operator over the relation to keep track of the view desc;
    // 2. Wrap the logical plan in a [[SubqueryAlias]] which tracks the name of the view.
    SubqueryAlias(multiParts, fromCatalogTable(metadata, isTempView = false))
 } else {
    // 物理表，生成 UnresolvedCatalogRelation
    SubqueryAlias(multiParts, UnresolvedCatalogRelation(metadata, options))
 }
}

private def getTempViewPlan(viewInfo: TemporaryViewRelation): View = viewInfo.plan match {
  case Some(p) => View(desc = viewInfo.tableMeta, isTempView = true, child = p)
  case None => fromCatalogTable(viewInfo.tableMeta, isTempView = true)
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
protected def analyzer: Analyzer = new Analyzer(catalog, conf) {
  override val extendedResolutionRules: Seq[Rule[LogicalPlan]] =
    new FindDataSourceTable(session) +:
      new ResolveSQLOnFile(session) +:
      customResolutionRules

  override val postHocResolutionRules: Seq[Rule[LogicalPlan]] =
    PreprocessTableCreation(session) +:
      PreprocessTableInsertion(conf) +:
      DataSourceAnalysis(conf) +:
      customPostHocResolutionRules

  override val extendedCheckRules: Seq[LogicalPlan => Unit] =
    PreWriteCheck +:
      PreReadCheck +:
      HiveOnlyCheck +:
      customCheckRules
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
/**
 * A logical query plan `Analyzer` with rules specific to Hive.
 */
override protected def analyzer: Analyzer = new Analyzer(catalog, conf) {
  override val extendedResolutionRules: Seq[Rule[LogicalPlan]] =
      // Determine the database, serde/format and schema of the Hive serde table, according to the storage properties.
      new ResolveHiveSerdeTable(session) +:
      // Replaces [[UnresolvedCatalogRelation]] with concrete relation logical plans.
      new FindDataSourceTable(session) +:
      new ResolveSQLOnFile(session) +:
      customResolutionRules

  override val postHocResolutionRules: Seq[Rule[LogicalPlan]] =
    new DetermineTableStats(session) +:
      RelationConversions(conf, catalog) +:
      PreprocessTableCreation(session) +:
      PreprocessTableInsertion(conf) +:
      DataSourceAnalysis(conf) +:
      HiveAnalysis +:
      customPostHocResolutionRules

  override val extendedCheckRules: Seq[LogicalPlan => Unit] =
    PreWriteCheck +:
      PreReadCheck +:
      customCheckRules
}

1
2
3
4
5
6
7
8
9
10
11
/**
 * Planner that takes into account Hive-specific strategies.
 */
override protected def planner: SparkPlanner = {
  new SparkPlanner(session.sparkContext, conf, experimentalMethods) with HiveStrategies {
    override val sparkSession: SparkSession = session

    override def extraPlanningStrategies: Seq[Strategy] =
      super.extraPlanningStrategies ++ customPlanningStrategies ++ Seq(HiveTableScans, Scripts)
 }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
// TableProvider
public interface TableProvider {
  StructType inferSchema(CaseInsensitiveStringMap options);
  default Transform[] inferPartitioning(CaseInsensitiveStringMap options) {
    return new Transform[0];
  }
  // getTable
  Table getTable(StructType schema, Transform[] partitioning, Map<String, String> properties);
  default boolean supportsExternalMetadata() {
    return false;
  }
}

// Table
public interface Table {
  String name();
  StructType schema();
  default Transform[] partitioning() {
    return new Transform[0];
  }
  default Map<String, String> properties() {
    return Collections.emptyMap();
  }
  Set<TableCapability> capabilities();
}

// SupportsRead
public interface SupportsRead extends Table {
  ScanBuilder newScanBuilder(CaseInsensitiveStringMap options);
}

// ScanBuilder
public interface ScanBuilder {
  Scan build();
}

// Scan
public interface Scan {
  StructType readSchema();
  default String description() {
    return this.getClass().toString();
  }
  default Batch toBatch() {
    throw new UnsupportedOperationException(description() + ": Batch scan are not supported");
  }
  default MicroBatchStream toMicroBatchStream(String checkpointLocation) {
    throw new UnsupportedOperationException(description() + ": Micro-batch scan are not supported");
  }
  default ContinuousStream toContinuousStream(String checkpointLocation) {
    throw new UnsupportedOperationException(description() + ": Continuous scan are not supported");
  }
  default CustomMetric[] supportedCustomMetrics() {
    return new CustomMetric[]{};
  }
}

1
2
3
4
5
6
spark.conf.set("spark.sql.catalog.mysql", classOf[JDBCTableCatalog].getName)
spark.conf.set("spark.sql.catalog.mysql.url", "jdbc:mysql://127.0.0.1/ttt?useSSL=false")
spark.conf.set("spark.sql.catalog.mysql.user", "work")
spark.conf.set("spark.sql.catalog.mysql.password", "*Work123")

val df_t1 = spark.sql("select * from mysql.ttt.t1")