fix: generate TabularSHAP samples lazily

core/src/main/scala/com/microsoft/azure/synapse/ml/explainers/TabularSHAP.scala

@@ -3,14 +3,14 @@
 
 package com.microsoft.azure.synapse.ml.explainers
 
-import breeze.stats.distributions.RandBasis
 import com.microsoft.azure.synapse.ml.core.schema.DatasetExtensions
 import com.microsoft.azure.synapse.ml.logging.FeatureNames
-import org.apache.spark.injections.UDFUtils
 import org.apache.spark.ml.ComplexParamsReadable
+import org.apache.spark.ml.linalg.SQLDataTypes.VectorType
 import org.apache.spark.ml.param.shared.HasInputCols
 import org.apache.spark.ml.util.Identifiable
-import org.apache.spark.sql.functions._
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.functions.{col, struct}
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.{DataFrame, Row}
 
@@ -53,30 +53,26 @@ class TabularSHAP(override val uid: String)
     })
 
     val infWeightVal = this.getInfWeight
-    val samplesUdf = UDFUtils.oldUdf(
-      {
-        (instance: Row, background: Row) =>
-          val sampler = new KernelSHAPTabularSampler(instance, background, effectiveNumSamples, infWeightVal)
-          (1 to effectiveNumSamples) map {
-            _ =>
-              implicit val randBasis: RandBasis = RandBasis.mt0
-              sampler.sample
-          }
-      },
-      getSampleSchema(sampleType)
+    val pairs = instances.crossJoin(background)
+    val outputSchema = StructType(
+      Seq(df.schema(idCol)) ++
+        sampleType.fields.map(field => field.copy(
+          dataType = TabularSHAP.nullableType(field.dataType), nullable = true)) ++
+        Seq(
+          StructField(coalitionCol, VectorType),
+          StructField(weightCol, DoubleType),
+          df.schema(targetClassesCol)
+        )
     )
+    implicit val outputEncoder: ExpressionEncoder[Row] = ExpressionEncoder(outputSchema)
 
-    val samplesCol = DatasetExtensions.findUnusedColumnName("samples", df)
-
-    instances.crossJoin(background)
-      .withColumn(samplesCol, explode(samplesUdf(col(instanceCol), col(backgroundCol))))
-      .select(
-        col(idCol),
-        expr(s"$samplesCol.$sampleField.*"),
-        col(samplesCol).getField(coalitionField).alias(coalitionCol),
-        col(samplesCol).getField(weightField).alias(weightCol),
-        col(targetClassesCol)
-      )
+    pairs.mapPartitions {
+      rows =>
+        TabularSHAP.expandPairs(rows, effectiveNumSamples) {
+          (instance, background) =>
+            new KernelSHAPTabularSampler(instance, background, effectiveNumSamples, infWeightVal)
+        }
+    }
   }
 
   override def validateSchema(schema: StructType): Unit = {
@@ -98,4 +94,33 @@ class TabularSHAP(override val uid: String)
   }
 }
 
-object TabularSHAP extends ComplexParamsReadable[TabularSHAP]
+object TabularSHAP extends ComplexParamsReadable[TabularSHAP] {
+  private def nullableType(dataType: DataType): DataType = dataType match {
+    case struct: StructType =>
+      StructType(struct.fields.map(field => field.copy(dataType = nullableType(field.dataType), nullable = true)))
+    case ArrayType(elementType, _) => ArrayType(nullableType(elementType), containsNull = true)
+    case MapType(keyType, valueType, _) =>
+      MapType(nullableType(keyType), nullableType(valueType), valueContainsNull = true)
+    case other => other
+  }
+
+  private[explainers] def sampleIterator[T](numSamples: Int)(sample: => T): Iterator[T] = {
+    Iterator.fill(numSamples)(()).map(_ => sample)
+  }
+
+  /** Expands (id, targetClasses, instance, background) pair rows into sample rows one at a time. */
+  private[explainers] def expandPairs(pairs: Iterator[Row], numSamples: Int)
+                                     (newSampler: (Row, Row) => KernelSHAPSampler[Row]): Iterator[Row] = {
+    pairs.flatMap {
+      row =>
+        val id = row.get(0)
+        val targetClasses = row.get(1)
+        val sampler = newSampler(row.getAs[Row](2), row.getAs[Row](3))
+
+        sampleIterator(numSamples)(sampler.sample).map {
+          case (sample, coalition, weight) =>
+            Row.fromSeq(Seq(id) ++ sample.toSeq ++ Seq(coalition, weight, targetClasses))
+        }
+    }
+  }
+}

core/src/test/scala/com/microsoft/azure/synapse/ml/explainers/split1/TabularSHAPExplainerSuite.scala

@@ -7,16 +7,19 @@ import com.microsoft.azure.synapse.ml.core.test.base.TestBase
 import com.microsoft.azure.synapse.ml.core.test.fuzzing.{TestObject, TransformerFuzzing}
 import com.microsoft.azure.synapse.ml.core.utils.BreezeUtils._
 import com.microsoft.azure.synapse.ml.explainers.LocalExplainer.KernelSHAP
-import com.microsoft.azure.synapse.ml.explainers.TabularSHAP
+import com.microsoft.azure.synapse.ml.explainers.{KernelSHAPTabularSampler, TabularSHAP}
 import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
-import org.apache.spark.ml.feature.{OneHotEncoder, StringIndexer, VectorAssembler}
+import org.apache.spark.ml.feature.{OneHotEncoder, SQLTransformer, StringIndexer, VectorAssembler}
 import org.apache.spark.ml.linalg.Vector
 import org.apache.spark.ml.util.MLReadable
 import org.apache.spark.ml.{Pipeline, PipelineModel}
-import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.functions.avg
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.functions.{array, avg, broadcast, lit}
+import org.apache.spark.sql.types._
 import org.scalactic.{Equality, TolerantNumerics}
 
+import java.util.concurrent.atomic.AtomicLong
+
 class TabularSHAPExplainerSuite extends TestBase
   with TransformerFuzzing[TabularSHAP] {
 
@@ -88,7 +91,158 @@ class TabularSHAPExplainerSuite extends TestBase
     assert(r2(0) === 1d)
   }
 
+  test("TabularSHAP sample expansion emits rows without buffering a pair's sample budget") {
+    val numFeatures = 14
+    val numSamples = 10000
+    var pairsPulled = 0
+    var samplesDrawn = 0
+    val instance = Row.fromSeq(Seq.fill(numFeatures)(1d))
+    val background = Row.fromSeq(Seq.fill(numFeatures)(0d))
+    val pairs = Iterator(0L, 1L).map {
+      id =>
+        pairsPulled += 1
+        Row(id, Seq(1), instance, background)
+    }
+
+    val rows = TabularSHAP.expandPairs(pairs, numSamples) {
+      (i, b) =>
+        new KernelSHAPTabularSampler(i, b, numSamples, 1E8) {
+          override def sample: (Row, Vector, Double) = {
+            samplesDrawn += 1
+            super.sample
+          }
+        }
+    }
+
+    assert(samplesDrawn === 0)
+    val first = rows.next()
+    assert(samplesDrawn === 1)
+    assert(pairsPulled === 1)
+    assert(first.length === numFeatures + 4)
+    assert(first.getLong(0) === 0L)
+    assert(rows.take(9).size === 9)
+    assert(samplesDrawn === 10)
+    assert(pairsPulled === 1)
+  }
+
+  test("TabularKernelSHAP accepts nullable background values") {
+    val instances = spark.range(3).selectExpr(
+      "cast(id as double) as x",
+      "cast(id * 2 as double) as y"
+    )
+    val background = Seq(
+      (Option.empty[Double], Option(1d)),
+      (Option(2d), Option.empty[Double])
+    ).toDF("x", "y")
+    val additiveModel = new SQLTransformer().setStatement(
+      "SELECT *, cast(coalesce(x, 10.0) + coalesce(y, 100.0) AS DOUBLE) AS score FROM __THIS__"
+    )
+
+    val results = KernelSHAP.tabular
+      .setInputCols(Array("x", "y"))
+      .setOutputCol("shapValues")
+      .setBackgroundData(background)
+      .setModel(additiveModel)
+      .setTargetCol("score")
+      .transform(instances)
+      .orderBy("x")
+      .select("shapValues")
+      .as[Seq[Vector]]
+      .collect()
+
+    assert(results.length === 3)
+    results.zipWithIndex.foreach {
+      case (Seq(values), x) =>
+        assert(values(0) === 56.5)
+        assert(values(1) === x.toDouble - 6d)
+        assert(values(2) === 2d * x - 50.5)
+    }
+  }
+
+  test("TabularSHAP createSamples streams encoded rows before expanding the sample budget") {
+    import TabularSHAPExplainerSuite._
+    val numSamples = 10000
+    val features = (0 until 14).map(i => s"f$i").toArray
+    val schema = StructType(Seq(StructField(features.head, new CountedValueUDT)) ++
+      features.tail.map(StructField(_, DoubleType)))
+    val instance = Row.fromSeq(Seq(CountedValue(1)) ++ Seq.fill(13)(1d))
+    val backgroundRow = Row.fromSeq(Seq(CountedValue(0)) ++ Seq.fill(13)(0d))
+    val instances = spark.createDataFrame(java.util.Collections.singletonList(instance), schema)
+      .withColumn("id", lit(0L)).withColumn("targetClasses", array(lit(1)))
+    val background = spark.createDataFrame(java.util.Arrays.asList(backgroundRow, backgroundRow), schema)
+    val samples = new SampleAccess()
+      .setInputCols(features).setNumSamples(numSamples).setBackgroundData(broadcast(background))
+      .samples(instances)
+
+    Serialized.set(0)
+    val rows = samples.limit(1).collect()
+    assert(rows.length === 1)
+    assert(rows.head.getAs[CountedValue]("f0").value === 0)
+    assert(rows.head.getAs[Vector]("coalition").toArray.forall(_ == 0d))
+    assert(rows.head.getAs[Double]("weight") === 1E8)
+    assert(Serialized.get() > 0 && Serialized.get() < numSamples,
+      s"Encoded ${Serialized.get()} feature values before returning the first sample")
+  }
+
+  test("TabularKernelSHAP preserves nulls inside background features") {
+    val featureShapes = Seq(
+      ("named_struct('value', %s)", "x.value"),
+      ("array(%s)", "x[0]"),
+      ("map('value', %s)", "x['value']"),
+      ("map('outer', array(named_struct('value', %s)))", "x['outer'][0].value")
+    )
+    featureShapes.foreach {
+      case (shape, access) =>
+        val instances = spark.range(1, 2).selectExpr(
+          s"${shape.format("cast(id as double)")} as x", "cast(id * 2 as double) as y")
+        val background = spark.range(1).selectExpr(
+          s"${shape.format("cast(null as double)")} as x", "cast(id as double) as y")
+        val additiveModel = new SQLTransformer().setStatement(
+          s"SELECT *, cast(coalesce($access, 10.0) + y AS DOUBLE) AS score FROM __THIS__")
+
+        val results = KernelSHAP.tabular
+          .setInputCols(Array("x", "y"))
+          .setOutputCol("shapValues")
+          .setBackgroundData(background)
+          .setModel(additiveModel)
+          .setTargetCol("score")
+          .transform(instances)
+        try {
+          val values = results.select("shapValues").as[Seq[Vector]].collect()
+          assert(values.length === 1)
+          assert(values.head.length === 1)
+          assert(values.head.head(0) === 10d)
+          assert(values.head.head(1) === -9d)
+          assert(values.head.head(2) === 2d)
+        } finally {
+          spark.catalog.clearCache()
+        }
+    }
+  }
+
   override def testObjects(): Seq[TestObject[TabularSHAP]] = Seq(new TestObject(kernelShap, infer))
 
   override def reader: MLReadable[_] = TabularSHAP
 }
+
+object TabularSHAPExplainerSuite {
+  // The local Spark test counts encoding at the production sampling boundary, not just iterator calls.
+  private val Serialized = new AtomicLong()
+
+  @SQLUserDefinedType(udt = classOf[CountedValueUDT])
+  case class CountedValue(value: Int)
+
+  class CountedValueUDT extends UserDefinedType[CountedValue] {
+    override def sqlType: DataType = IntegerType
+    override def userClass: Class[CountedValue] = classOf[CountedValue]
+    override def serialize(value: CountedValue): Any = {
+      Serialized.incrementAndGet()
+      value.value
+    }
+    override def deserialize(value: Any): CountedValue = CountedValue(value.asInstanceOf[Int])
+  }
+
+  private class SampleAccess extends TabularSHAP {
+    def samples(df: DataFrame): DataFrame = createSamples(df, "id", "coalition", "weight", "targetClasses")
+  }
+}

docs/Explore Algorithms/Responsible AI/Tabular Explainers.ipynb

@@ -164,7 +164,9 @@
     }
    },
    "source": [
-    "We create a TabularSHAP explainer, set the input columns to all the features the model takes, specify the model and the target output column we're trying to explain. In this case, we're trying to explain the \"probability\" output, which is a vector of length 2, and we're only looking at class 1 probability. Specify targetClasses to `[0, 1]` if you want to explain class 0 and 1 probability at the same time. Finally we sample 100 rows from the training data for background data, which is used for integrating out features in Kernel SHAP."
+    "We create a TabularSHAP explainer, set the input columns to all the features the model takes, specify the model and the target output column we're trying to explain. In this case, we're trying to explain the \"probability\" output, which is a vector of length 2, and we're only looking at class 1 probability. Specify targetClasses to `[0, 1]` if you want to explain class 0 and 1 probability at the same time. Finally we sample 100 rows from the training data for background data, which is used for integrating out features in Kernel SHAP.\n",
+    "\n",
+    "Sample rows are generated incrementally, but model-scoring work still scales with the number of explained rows, background rows, and effective samples per pair. Start with a small representative background and sample budget before scaling up. Background features can contain nulls, including nested nulls in arrays, maps, and structs, provided the model can score those values."
    ]
   },
   {