Dataframe processing generically using Scala

Question

This code below I understand and was helpful.

But I would like to make this a generic approach, but cannot actually get started, and think that it is not possible actually with the case statement. I am looking at another approach, but am interested if a generic approach is also possible here.

import spark.implicits._
import org.apache.spark.sql.Encoders

// Creating case classes with the schema of your json objects. We're making
// these to make use of strongly typed Datasets. Notice that the MyChgClass has
// each field as an Option: this will enable us to choose between "chg" and
// "before"
case class MyChgClass(b: Option[String], c: Option[String], d: Option[String])
case class MyFullClass(k: Int, b: String, c: String, d: String)
case class MyEndClass(id: Int, after: MyFullClass)

// Creating schemas for the from_json function
val chgSchema = Encoders.product[MyChgClass].schema
val beforeSchema = Encoders.product[MyFullClass].schema

// Your dataframe from the example
val df = Seq(
  (1, """{"b": "new", "c": "new"}""",  """{"k": 1, "b": "old", "c": "old", "d": "old"}""" ),
  (2, """{"b": "new", "d": "new"}""",  """{"k": 2, "b": "old", "c": "old", "d": "old"}""" )
).toDF("id", "chg", "before")

// Parsing the json string into our case classes and finishing by creating a
// strongly typed dataset with the .as[] method
val parsedDf = df
  .withColumn("parsedChg",from_json(col("chg"), chgSchema))
  .withColumn("parsedBefore",from_json(col("before"), beforeSchema))
  .drop("chg")
  .drop("before")
  .as[(Int, MyChgClass, MyFullClass)]

// Mapping over our dataset with a lot of control of exactly what we want. Since
// the "chg" fields are options, we can use the getOrElse method to choose
// between either the "chg" field or the "before" field
val output = parsedDf.map{
  case (id, chg, before) => {
    MyEndClass(id, MyFullClass(
      before.k,
      chg.b.getOrElse(before.b),
      chg.c.getOrElse(before.c),
      chg.d.getOrElse(before.d)
    ))
  }
}

output.show(false)
parsedDf.printSchema()

We have many such situations, but with differing payload. I can get the fields of the case class, but cannot see the forest for the trees how to make this generic. E,g, [T] type approach for the below. I am wondering if this can be done in fact?

I can get a List of attributes, and am wondering if something like attrList.map(x => ...) with substitution can be used for the chg.b etc?

  val output = parsedDf.map{
      case (id, chg, before) => {
        MyEndClass(id, MyFullClass(
          before.k,
          chg.b.getOrElse(before.b),
          chg.c.getOrElse(before.c),
          chg.d.getOrElse(before.d)
        ))
      }
    }

Answer 1

Does the following macro work for your use case?

// libraryDependencies += scalaOrganization.value % "scala-reflect" % scalaVersion.value
import scala.language.experimental.macros
import scala.reflect.macros.blackbox

def mkInstance[A, B](before: A, chg: B): A = macro mkInstanceImpl[A]

def mkInstanceImpl[A: c.WeakTypeTag](c: blackbox.Context)(before: c.Tree, chg: c.Tree): c.Tree = {
  import c.universe._

  val A = weakTypeOf[A]

  val classAccessors = A.decls.collect {
    case m: MethodSymbol if m.isCaseAccessor => m
  }

  val arg = q"$before.${classAccessors.head}"
  val args = classAccessors.tail.map(m => q"$chg.${m.name}.getOrElse($before.$m)")

  q"new $A($arg, ..$args)"
}

// in a different subproject

val output = parsedDf.map{
  case (id, chg, before) => {
    MyEndClass(id,
      mkInstance(before, chg)
    )
  }
}

    // scalacOptions += "-Ymacro-debug-lite"
// scalac: new MyFullClass(before.k, chg.b.getOrElse(before.b), chg.c.getOrElse(before.c), chg.d.getOrElse(before.d))

https://scastie.scala-lang.org/bXq5FHb3QuC5PqlhZOfiqA

Alternatively you can use Shapeless

// libraryDependencies += "com.chuusai" %% "shapeless" % "2.3.10"
import shapeless.{Generic, HList, LabelledGeneric, Poly2}
import shapeless.ops.hlist.{IsHCons, Mapped, ZipWith}
import shapeless.ops.record.Keys

def mkInstance[A, B, L <: HList, H, T <: HList, OptT <: HList, L1 <: HList, T1 <: HList, T2 <: HList, K <: HList](
  before: A, chg: B
)(implicit
  // checking that field names in tail of A are equal to field names in B  
  aLabelledGeneric: LabelledGeneric.Aux[A, L1],
  bLabelledGeneric: LabelledGeneric.Aux[B, T2],
  isHCons1: IsHCons.Aux[L1, _, T1],
  keys: Keys.Aux[T1, K],
  keys1: Keys.Aux[T2, K],
  // checking that field types in B are Options of field types in tail of A 
  aGeneric: Generic.Aux[A, L],
  isHCons: IsHCons.Aux[L, H, T],
  mapped: Mapped.Aux[T, Option, OptT],
  bGeneric: Generic.Aux[B, OptT],
  zipWith: ZipWith.Aux[OptT, T, getOrElsePoly.type, T],
): A = {
  val aHList = aGeneric.to(before)
  aGeneric.from(isHCons.cons(isHCons.head(aHList), zipWith(bGeneric.to(chg), isHCons.tail(aHList))))
}

object getOrElsePoly extends Poly2 {
  implicit def cse[A]: Case.Aux[Option[A], A, A] = at(_ getOrElse _)
}

Since all the classes are now known at compile-time it's better to use compile-time reflection (macros themselves or macros hidden inside type classes as in Shapeless) but in principle runtime reflection also can be used

import scala.reflect.ClassTag
import scala.reflect.runtime.{currentMirror => rm}
import scala.reflect.runtime.universe._

def mkInstance[A: TypeTag : ClassTag, B: TypeTag : ClassTag](before: A, chg: B): A = {
  val A = typeOf[A]
  val B = typeOf[B]
  val classAccessors = A.decls.collect {
    case m: MethodSymbol if m.isCaseAccessor => m
  }.toList
  val arg = rm.reflect(before).reflectMethod(classAccessors.head)()
  val args = classAccessors.tail.map(m =>
    rm.reflect(chg).reflectMethod(B.decl(m.name).asMethod)()
      .asInstanceOf[Option[_]].getOrElse(
        rm.reflect(before).reflectMethod(m)()
      )
  )
  rm.reflectClass(A.typeSymbol.asClass)
    .reflectConstructor(A.decl(termNames.CONSTRUCTOR).asMethod)(arg :: args : _*)
    .asInstanceOf[A]
}

Answer 2

Although title of question is a little different, I asked the question on how to process data generically from a feed of what is ostensibly a CDC feed. This approach handles the above also, plus some schema evolution aspects. Hence, it may be of interest to others.

A simpler approach. Use compile-time reflection, and allowing for schema changes, a simple approach for JSON data from CDC via KAFKA can be done like this, shown for two types of records from Shareplex, in the same feed - you need to adapt to your own situation:

...
import org.json4s._
import org.json4s.jackson.JsonMethods._
import spark.implicits._
import org.apache.spark.sql.functions.{col, lit, when, from_json, map_keys, map_values, regexp_replace, coalesce}
import org.apache.spark.sql.types.{MapType, StringType}
import org.apache.spark.sql.Encoders
import org.apache.spark.sql.types.{MapType, StringType, StructType, IntegerType}

case class MyMeta(op: String, table: String)
val metaSchema = Encoders.product[MyMeta].schema
case class MySales(NUM: Option[Integer], PRODUCT_ID: Option[String], DESCRIPTION: Option[String], OLD_FIELD_1: Option[String]) 
val salesSchema = Encoders.product[MySales].schema
case class MyProducts(PRODUCT_ID: Option[String], DESCRIPTION: Option[String], PRICE: Option[Int], OLD_FIELD_1: Option[String]) 
val ProductsSchema = Encoders.product[MyProducts].schema

def getAfterImage (op: String, data: String, key: String, jsonOLD_TABLE_FIELDS: String) : String = {   
  val jsonOLD_FIELDS = parse(jsonOLD_TABLE_FIELDS)   
  val jsonData = parse(data)                         
  val jsonKey = parse(key)                           
   
  op match {
  case "ins" =>
               return(compact(render(jsonData merge jsonOLD_FIELDS)))
  case _ => 
               val Diff(changed, added, deleted) = jsonKey diff jsonData
               return(compact(render(changed merge deleted merge jsonOLD_FIELDS)))
  }
}
val afterImage = spark.udf.register("callUDFAI", getAfterImage _)

val path = "/FileStore/tables/json_DE_CDC_SHAREPLEX_file.txt"  
val df = spark.read.text(path)  // String.
val df2 = df.withColumn("value", from_json(col("value"), MapType(StringType, StringType)))    
val df3 = df2.select(map_values(col("value")))  
val df4 = df3.select($"map_values(value)"(0).as("meta"), $"map_values(value)"(1).as("data"), $"map_values(value)"(2).as("key")).withColumn("parsedMeta", from_json(col("meta"), metaSchema)).drop("meta").select(col("parsedMeta.*"), col("data"), col("key")).withColumn("key2", coalesce(col("key"), lit(""" { "DUMMY_FIELD_XXX": ""} """) )).toDF().cache()

val df_sales    = df4.filter('table === "BILL.SALES") 
val df_products = df4.filter('table === "BILL.PRODUCTS")
val df_sales_output = df_sales.withColumn("afterImage", afterImage(col("op"), col("data"), col("key2") , lit(""" { "OLD_FIELD_1": ""} """)))
                              .select("afterImage") 
val df_products_output = df_products.withColumn("afterImage", afterImage(col("op"), col("data"), col("key2") , lit(""" { "OLD_FIELD_A":"", "OLD_FIELD_B":""} """)))
                                    .select("afterImage")                          
val df_sales_output_final = df_sales_output.withColumn("parsedSales", from_json(col("afterImage"), salesSchema)) 
...