How can I introspect and pre-load all collections from MongoDB into the Spark SQL catalog?

user_304151 0
Viewed 98

When learning Spark SQL, I've been using the following approach to register a collection into the Spark SQL catalog and query it.

val persons: Seq[MongoPerson] = Seq(MongoPerson("John", "Doe"))

sqlContext.createDataset(persons)
  .write
  .format("com.mongodb.spark.sql.DefaultSource")
  .option("collection", "peeps")
  .mode("append")
  .save()

sqlContext.read
  .format("com.mongodb.spark.sql.DefaultSource")
  .option("collection", "peeps")
  .load()
  .as[Peeps]
  .show()

However, when querying it, it seems that I need to register it as a temporary view in order to access it using SparkSQL.

val readConfig = ReadConfig(Map("uri" -> "mongodb://localhost:37017/test", "collection" -> "morepeeps"), Some(ReadConfig(spark)))
val people: DataFrame = MongoSpark.load[Peeps](spark, readConfig)
people.show()
people.createOrReplaceTempView("peeps")

spark.catalog.listDatabases().show()
spark.catalog.listTables().show()

sqlContext.sql("SELECT * FROM peeps")
  .as[Peeps]
  .show()

For a database with quite a few collections, is there a way to hydrate the Spark SQL schema catalog so that this op isn't so verbose?

apache-spark-sql apache-spark mongodb scala
1 Answers

So there's a couple things going on. First of all, simply loading the Dataset using sqlContext.read will not register it with SparkSQL catalog. The end of the function chain you have in your first code sample returns a Dataset at .as[Peeps]. You need to tell Spark that you want to use it as a view.

Depending on what you're doing with it, I might recommend leaning on the Scala Dataset API rather than SparkSQL. However, if SparkSQL is absolutely essential, you can likely speed things up programmatically.

In my experience, you'll need to run that boilerplate on each table you want to import. Fortunately, Scala is a proper programming language, so we can cut down on code duplication substantially by using a function, and calling it as such:

val MongoDbUri: String = "mongodb://localhost:37017/test" // store this as a constant somewhere

// T must be passed in as some case class
// Note, you can also add a second parameter to change the view name if so desired
def loadTableAsView[T <: Product : TypeTag](table: String)(implicit spark: SparkSession): Dataset[T] {
  val configMap = Map(
    "uri" -> MongoDbUri,
    "collection" -> table
  )
  val readConfig = ReadConfig(configMap, Some(ReadConfig(spark)))
  val df: DataFrame = MongoSpark.load[T](spark, readConfig)
  df.createOrReplaceTempView(table)
  df.as[T]
}

And to call it:

// Note: if spark is defined implicitly, e.g. implicit val spark: SparkSession = spark, you won't need to pass it explicitly

val peepsDS: Dataset[Peeps] = loadTableAsView[Peeps]("peeps")(spark)
val chocolatesDS: Dataset[Chocolates] = loadTableAsView[Chocolates]("chocolates")(spark)
val candiesDS: Dataset[Candies] = loadTableAsView[Candies]("candies")(spark)

spark.catalog.listDatabases().show()
spark.catalog.listTables().show()

peepsDS.show()
chocolatesDS.show()
candiesDS.show()

This will substantially cut down your boilerplate, and also allow you to more easily write some tests for that repeated bit of code. There's also probably a way to create a map of table names to case classes that you can then iterate over, but I don't have an IDE handy to test it out.

user_104306680
Related