comments

Author: Sophie Wang

api/src/main/scala/ai/chronon/api/Constants.scala

@@ -39,6 +39,6 @@ object Constants {
   val ContextualSourceValues: String = "contextual_values"
   val TeamOverride: String = "team_override"
   val LabelColumnPrefix: String = "label"
-  val LabelViewPropertyFeatureTable: String = "featureTable"
-  val LabelViewPropertyKeyLabelTable: String = "labelTable"
+  val LabelViewPropertyFeatureTable: String = "feature_table"
+  val LabelViewPropertyKeyLabelTable: String = "label_table"
 }

api/src/main/scala/ai/chronon/api/Extensions.scala

@@ -8,6 +8,7 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import java.io.{PrintWriter, StringWriter}
 import java.util
 import scala.collection.mutable
+import scala.collection.JavaConverters._
 import scala.util.{Failure, ScalaVersionSpecificCollectionsConverter, Success, Try}
 
 object Extensions {
@@ -567,9 +568,12 @@ object Extensions {
         .distinct
     }
 
-    // a list of columns which can identify a row on left
-    def rowIdentifier: Array[String] = {
-      leftKeyCols ++ Array(Constants.PartitionColumn)
+    // a list of columns which can identify a row on left, use user specified columns by default
+    def rowIdentifier(userRowId: util.List[String] = null): Array[String] = {
+      if (userRowId != null && !userRowId.isEmpty)
+        userRowId.asScala.toArray
+      else
+        leftKeyCols ++ Array(Constants.PartitionColumn)
     }
   }
 

spark/src/main/scala/ai/chronon/spark/JoinUtils.scala

@@ -125,7 +125,7 @@ object JoinUtils {
       tableUtils.getSchemaFromTable(rightTable)
         .filterNot(field => joinKeys.contains(field.name))
         .map(field => {
-          if(field.name.startsWith("label")) {
+          if(field.name.startsWith(labelColumnPrefix)) {
             s"r.${field.name}"
           } else {
             s"r.${field.name} AS ${labelColumnPrefix}_${field.name}"
@@ -199,7 +199,7 @@ object JoinUtils {
 
   /**
    * compute the mapping label_ds -> PartitionRange of ds which has this label_ds as latest version
-   *  - Get all partitions from table and
+   *  - Get all partitions from table
    *  - For each ds, find the latest available label_ds
    *  - Reverse the mapping and get the ds partition range for each label version(label_ds)
    *

spark/src/main/scala/ai/chronon/spark/LabelJoin.scala

@@ -59,12 +59,10 @@ class LabelJoin(joinConf: api.Join, tableUtils: TableUtils, labelDS: String) {
     } else {
       // creating final join view with feature join output table
       println(s"Joining label table : ${outputLabelTable} with joined output table : ${joinConf.metaData.outputTable}")
-      val joinKeys: Array[String] = if (joinConf.rowIds != null && !joinConf.rowIds.isEmpty)
-        joinConf.rowIds.asScala.toArray else labelJoinConf.rowIdentifier
       JoinUtils.createOrReplaceView(joinConf.metaData.outputFinalView,
                                     leftTable = joinConf.metaData.outputTable,
                                     rightTable = outputLabelTable,
-                                    joinKeys = joinKeys,
+                                    joinKeys = labelJoinConf.rowIdentifier(joinConf.rowIds),
                                     tableUtils = tableUtils,
                                     viewProperties = Map(Constants.LabelViewPropertyKeyLabelTable -> outputLabelTable,
                                       Constants.LabelViewPropertyFeatureTable -> joinConf.metaData.outputTable))
@@ -128,7 +126,7 @@ class LabelJoin(joinConf: api.Join, tableUtils: TableUtils, labelDS: String) {
     finalResult
   }
 
-  def computeRange(leftDf: DataFrame, leftRange: PartitionRange, labelDs: String): DataFrame = {
+  def computeRange(leftDf: DataFrame, leftRange: PartitionRange, sanitizedLabelDs: String): DataFrame = {
     val leftDfCount = leftDf.count()
     val leftBlooms = labelJoinConf.leftKeyCols.par.map { key =>
       key -> leftDf.generateBloomFilter(key, leftDfCount, joinConf.left.table, leftRange)
@@ -145,7 +143,7 @@ class LabelJoin(joinConf: api.Join, tableUtils: TableUtils, labelDS: String) {
       }
     }
 
-    val rowIdentifier = labelJoinConf.rowIdentifier
+    val rowIdentifier = labelJoinConf.rowIdentifier(joinConf.rowIds)
     println("Label Join filtering left df with only row identifier:", rowIdentifier.mkString(", "))
     val leftFiltered = JoinUtils.filterColumns(leftDf, rowIdentifier)
 
@@ -154,7 +152,7 @@ class LabelJoin(joinConf: api.Join, tableUtils: TableUtils, labelDS: String) {
     }
 
     // assign label ds value and drop duplicates
-    val updatedJoin = joined.withColumn(Constants.LabelPartitionColumn, lit(labelDS))
+    val updatedJoin = joined.withColumn(Constants.LabelPartitionColumn, lit(sanitizedLabelDs))
       .dropDuplicates(rowIdentifier)
     updatedJoin.explain()
     updatedJoin.drop(Constants.TimePartitionColumn)

spark/src/test/scala/ai/chronon/spark/test/FeatureWithLabelJoinTest.scala

@@ -36,26 +36,22 @@ class FeatureWithLabelJoinTest {
     val runner = new LabelJoin(joinConf, tableUtils, labelDS)
     val labelDf = runner.computeLabelJoin()
     println(" == First Run Label version 2022-10-30 == ")
-    prefixColumnName(labelDf, exceptions = labelJoinConf.rowIdentifier).show()
+    prefixColumnName(labelDf, exceptions = labelJoinConf.rowIdentifier()).show()
     val featureDf = tableUtils.sparkSession.table(joinConf.metaData.outputTable)
     println(" == Features == ")
     featureDf.show()
     val computed = tableUtils.sql(s"select * from ${joinConf.metaData.outputFinalView}")
     computed.show()
-    val expectedFinal = featureDf.join(prefixColumnName(labelDf, exceptions = labelJoinConf.rowIdentifier),
-      labelJoinConf.rowIdentifier,
+    val expectedFinal = featureDf.join(prefixColumnName(labelDf, exceptions = labelJoinConf.rowIdentifier()),
+      labelJoinConf.rowIdentifier(),
       "left_outer")
     println(" == Expected == ")
     expectedFinal.show()
     val diff = Comparison.sideBySide(computed,
       expectedFinal,
       List("listing",
         "ds",
-        "label_ds",
-        "feature_review",
-        "feature_locale",
-        "label_listing_labels_dim_bedrooms",
-        "label_listing_labels_dim_room_type"))
+        "label_ds"))
     if (diff.count() > 0) {
       println(s"Actual count: ${computed.count()}")
       println(s"Expected count: ${expectedFinal.count()}")
@@ -98,7 +94,7 @@ class FeatureWithLabelJoinTest {
     println(exceptions.mkString(", "))
     val renamedColumns = df.columns
       .map(col => {
-        if(exceptions.contains(col) || col.startsWith("label")) {
+        if(exceptions.contains(col) || col.startsWith(prefix)) {
           df(col)
         } else {
           df(col).as(s"$prefix$col")

spark/src/test/scala/ai/chronon/spark/test/JoinUtilsTest.scala

@@ -254,7 +254,8 @@ class JoinUtilsTest {
     val keys = Array("listing_id", Constants.PartitionColumn)
 
     JoinUtils.createOrReplaceView(finalViewName, leftTableName, rightTableName, keys, tableUtils,
-      viewProperties = Map("featureTable" -> leftTableName, "labelTable" -> rightTableName))
+      viewProperties = Map(Constants.LabelViewPropertyFeatureTable -> leftTableName,
+                       Constants.LabelViewPropertyKeyLabelTable -> rightTableName))
     val view = tableUtils.sql(s"select * from $finalViewName")
     view.show()
     assertEquals(6, view.count())
@@ -269,10 +270,13 @@ class JoinUtilsTest {
     assertEquals(0, latest.filter(latest("listing_id") === "3").count())
     assertEquals("2022-11-22", latest.where(latest("ds") === "2022-10-07").
       select("label_ds").first().get(0))
+    // label_ds should be unique per ds + listing
+    val removeDup = latest.dropDuplicates(Seq("label_ds", "ds"))
+    assertEquals(removeDup.count(), latest.count())
 
     val properties = tableUtils.getTableProperties(latestLabelView)
     assertTrue(properties.isDefined)
-    assertEquals(properties.get.get("featureTable"), Some(leftTableName))
+    assertEquals(properties.get.get(Constants.LabelViewPropertyFeatureTable), Some(leftTableName))
     assertEquals(properties.get.get("newProperties"), Some("value"))
   }