fix: preserve VW Generic input schemas

vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala

@@ -11,7 +11,7 @@ import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.util.Identifiable
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession, functions => F}
-import org.apache.spark.sql.types.{StringType, StructField, StructType}
+import org.apache.spark.sql.types.{StringType, StructType}
 
 /**
   * This Estimator supports driving VW using VW-style examples (e.g. 0 |a b c)
@@ -82,6 +82,12 @@ class VowpalWabbitGeneric(override val uid: String)
 
 object VowpalWabbitGeneric extends ComplexParamsReadable[VowpalWabbitGeneric]
 
+/**
+  * Predicts from string examples while preserving every top-level input field in its original order,
+  * including its type, nullability and metadata. Existing flat prediction fields are appended.
+  * Input names that conflict with prediction names must be renamed before transforming.
+  * Name comparisons follow Spark's case-sensitivity setting; the default is case-insensitive.
+  */
 class VowpalWabbitGenericModel(override val uid: String)
   extends Model[VowpalWabbitGenericModel]
     with VowpalWabbitBaseModel
@@ -95,19 +101,40 @@ class VowpalWabbitGenericModel(override val uid: String)
 
   override def copy(extra: ParamMap): this.type = defaultCopy(extra)
 
-  private def schemaForPredictionType: StructType =
-    StructType(StructField(getInputCol, StringType, nullable = false) +: VowpalWabbitPrediction.getSchema(vw).fields)
+  // Preserve top-level input fields with their types, metadata and nullability.
+  // Prediction fields stay flat; callers must rename any conflicting input columns.
+  private def schemaForPredictionType(schema: StructType, caseSensitive: Boolean): StructType = {
+    require(schema(getInputCol).dataType == StringType,
+      s"Input column '$getInputCol' must have StringType.")
+    val predictionFields = VowpalWabbitPrediction.getSchema(vw).fields
+    val conflicts = schema.fieldNames.filter { inputName =>
+      predictionFields.exists { field =>
+        if (caseSensitive) inputName == field.name else inputName.equalsIgnoreCase(field.name)
+      }
+    }
+    require(conflicts.isEmpty,
+      s"Input columns conflict with prediction fields: ${conflicts.mkString(", ")}. Rename them before transforming.")
+    StructType(schema.fields ++ predictionFields)
+  }
 
+  /**
+    * Keeps all input fields with their type, nullability and metadata at every nesting level.
+    * Appends flat predictions.
+    * The input example column must be a string. Conflicting names are rejected using
+    * the dataset session's spark.sql.caseSensitive setting, including inputCol itself.
+    */
   override def transform(dataset: Dataset[_]): DataFrame = {
     // this is doing predict, but lightgbm also logs logTransform in the model...
     logTransform ({
       val df = dataset.toDF()
       val inputColIdx = df.schema.fieldIndex(getInputCol)
 
       val predictToSeq = VowpalWabbitPrediction.getPredictionFunc(vw)
-      val rowEncoder = ExpressionEncoder(schemaForPredictionType)
+      val caseSensitive = df.sparkSession.conf.get("spark.sql.caseSensitive", "false").trim.toBoolean
+      val outputSchema = schemaForPredictionType(df.schema, caseSensitive)
+      val rowEncoder = ExpressionEncoder(outputSchema)
 
-      df.mapPartitions(inputRows => {
+      val encoded = df.mapPartitions(inputRows => {
         inputRows.map { row => {
           val input = row.getString(inputColIdx)
           val prediction = vw.predictFromString(input)
@@ -116,15 +143,23 @@ class VowpalWabbitGenericModel(override val uid: String)
         }}
       })(rowEncoder)
         .toDF()
+
+      val internalColumns = outputSchema.fields.indices.map(index => s"__vw_generic_out_$index")
+      encoded.toDF(internalColumns: _*).select(outputSchema.fields.zip(internalColumns).map {
+        case (field, internalName) => F.col(internalName).cast(field.dataType).as(field.name, field.metadata)
+      }: _*)
     }, dataset.columns.length)
   }
 
+  /**
+    * Preserves all input fields and appends flat prediction fields, rejecting non-string examples
+    * and conflicting names. Uses the active Spark session's case-sensitivity setting, or false
+    * when no session is active; callers should use the same session as the dataset to transform.
+    */
   override def transformSchema(schema: StructType): StructType = {
-    // validate the inputCol is present
-    // TODO: is this the right way of doing it?
-    schema.fieldIndex(getInputCol)
-
-    StructType(schemaForPredictionType)
+    val caseSensitive = SparkSession.getActiveSession
+      .exists(_.conf.get("spark.sql.caseSensitive", "false").trim.toBoolean)
+    schemaForPredictionType(schema, caseSensitive)
   }
 }
 

vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala

@@ -5,8 +5,10 @@ package com.microsoft.azure.synapse.ml.vw
 
 import com.microsoft.azure.synapse.ml.build.BuildInfo
 import com.microsoft.azure.synapse.ml.core.env.FileUtilities
+import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using
 import com.microsoft.azure.synapse.ml.core.test.benchmarks.{Benchmarks}
 import com.microsoft.azure.synapse.ml.core.test.fuzzing.{EstimatorFuzzing, TestObject}
+import org.apache.spark.ml.feature.{IndexToString, StringIndexer}
 import org.apache.spark.ml.util.MLReadable
 import org.apache.spark.sql.expressions.Window
 import org.apache.spark.sql.{functions => F, types => T}
@@ -240,6 +242,133 @@ class VerifyVowpalWabbitGeneric extends Benchmarks with EstimatorFuzzing[VowpalW
     verifyResult(expected, actual)
   }
 
+  test("Generic model preserves complete input schemas for scalar and array predictions") {
+    val metadata = new T.MetadataBuilder().putString("source", "generic-schema-test").build()
+    val examples = Seq("1 |a b c", "2 |d e f").toDF("value").coalesce(1)
+    for (probabilities <- Seq(false, true)) {
+      val learner = new VowpalWabbitGeneric().setUseBarrierExecutionMode(false)
+      if (probabilities) learner.setPassThroughArgs("--oaa 2")
+      val model = learner.fit(examples)
+      if (probabilities) model.setTestArgs("--probabilities")
+      val control = model.transform(examples)
+      val predictionNames = control.columns.drop(1)
+      val input = examples.withColumn("id", F.length($"value"))
+        .withColumn("label", F.lit("retained"))
+        .withColumn("value", F.when($"value".isNotNull, $"value"))
+        .select($"id", $"value".as("value", metadata), $"label")
+      assert(input.schema("value").nullable)
+      for (columns <- Seq(Seq("value", "id", "label"), Seq("id", "value", "label"),
+        Seq("id", "label", "value"))) {
+        val arranged = input.select(columns.map(F.col): _*)
+        val transformed = model.transform(arranged)
+        assert(transformed.schema == model.transformSchema(arranged.schema))
+        assert(transformed.schema.fields.take(columns.length).toSeq == arranged.schema.fields.toSeq,
+          "GENERIC_INPUT_SCHEMA_PRESERVATION_REGRESSION")
+        assert(transformed.columns.drop(columns.length).toSeq == predictionNames.toSeq)
+        assertDFEq(arranged, transformed.select(columns.map(F.col): _*))
+        assertDFEq(control.orderBy("value").select(predictionNames.map(F.col): _*),
+          transformed.orderBy("value").select(predictionNames.map(F.col): _*))
+      }
+      val renamed = input.withColumnRenamed("value", "examples")
+      val custom = model.copy(org.apache.spark.ml.param.ParamMap.empty).setInputCol("examples")
+      assert(custom.transform(renamed).schema == custom.transformSchema(renamed.schema))
+      assertDFEq(model.transform(input).orderBy("value"),
+        custom.transform(renamed).withColumnRenamed("examples", "value").orderBy("value"))
+      intercept[IllegalArgumentException] { model.transformSchema(T.StructType(Seq.empty)) }
+      intercept[IllegalArgumentException] { model.transform(input.drop("value")) }
+      val nonString = input.withColumn("value", F.lit(1))
+      val typeError = intercept[IllegalArgumentException] { model.transformSchema(nonString.schema) }
+      assert(typeError.getMessage.contains("must have StringType"))
+      intercept[IllegalArgumentException] { model.transform(nonString) }
+      if (probabilities) {
+        val conflicting = input.withColumn("predictions", F.lit(1))
+        intercept[IllegalArgumentException] { model.transformSchema(conflicting.schema) }
+        intercept[IllegalArgumentException] { model.transform(conflicting) }
+      }
+    }
+  }
+
+  test("Generic model preserves nested field metadata and nullability") {
+    val metadata = new T.MetadataBuilder().putString("source", "nested-schema-test").build()
+    val examples = Seq("1 |a b c", "0 |d e f").toDF("value").coalesce(1)
+    val model = new VowpalWabbitGeneric().setUseBarrierExecutionMode(false).fit(examples)
+    using(model.vw) { _ =>
+      for (nullableChild <- Seq(false, true); nullableParent <- Seq(false, true)) {
+        val child = if (nullableChild) F.when($"value".startsWith("1"), F.lit(7)) else F.lit(7)
+        val nested = F.struct(child.as("child", metadata))
+        val nullStruct = F.when(F.lit(false), nested)
+        val payloads = Seq(
+          nested,
+          F.array(nested, nullStruct),
+          F.map(F.lit("key"), nested, F.lit("null"), nullStruct),
+          F.map(nested, nested))
+        for (payload <- payloads) {
+          val field = if (nullableParent) F.when($"value".startsWith("1"), payload) else payload
+          val input = examples.select($"value", field.as("payload", metadata))
+          val output = model.transform(input)
+          assert(output.schema == model.transformSchema(input.schema))
+          assert(output.schema.fields.take(input.columns.length).toSeq == input.schema.fields.toSeq)
+          assert(output.select("value", "payload").collect().toSeq == input.collect().toSeq)
+        }
+      }
+    }.get
+  }
+
+  test("Generic model preserves nested categorical metadata for downstream decoding") {
+    val examples = Seq(("1 |a b c", "red"), ("0 |d e f", "blue"))
+      .toDF("value", "category").coalesce(1)
+    val indexer = new StringIndexer().setInputCol("category").setOutputCol("indexed").fit(examples)
+    val input = indexer.transform(examples).select($"value", F.struct($"indexed").as("payload"))
+    val decoder = new IndexToString().setInputCol("indexed").setOutputCol("decoded")
+    val expected = decoder.transform(input.select("payload.indexed"))
+      .select("decoded").as[String].collect().toSeq
+    assert(expected == Seq("red", "blue"))
+    val model = new VowpalWabbitGeneric().setUseBarrierExecutionMode(false).fit(examples.select("value"))
+    using(model.vw) { _ =>
+      val output = model.transform(input)
+      val actual = decoder.transform(output.select("payload.indexed"))
+        .select("decoded").as[String].collect().toSeq
+      assert(actual == expected)
+    }.get
+  }
+
+  test("Generic prediction name collisions follow Spark case sensitivity") {
+    val examples = Seq("1 |a b c", "0 |d e f").toDF("value").coalesce(1)
+    val model = new VowpalWabbitGeneric().setUseBarrierExecutionMode(false).fit(examples)
+    val original = spark.conf.get("spark.sql.caseSensitive", "false")
+    try {
+      for (caseSensitive <- Seq(false, true)) {
+        spark.conf.set("spark.sql.caseSensitive", caseSensitive.toString)
+        for (name <- Seq("prediction", "confidence", "Prediction", "Confidence")) {
+          val input = examples.withColumn(name, F.lit(7))
+          val conflicts = !caseSensitive || name == name.toLowerCase(java.util.Locale.ROOT)
+          if (conflicts) {
+            val schemaError = intercept[IllegalArgumentException] { model.transformSchema(input.schema) }
+            assert(schemaError.getMessage.contains("Rename them before transforming"))
+            intercept[IllegalArgumentException] { model.transform(input) }
+          } else {
+            assert(model.transform(input).schema == model.transformSchema(input.schema))
+            assert(model.transform(input).collect().length == 2)
+          }
+        }
+        val nearMiss = examples.withColumn("prediction_input", F.lit(7))
+        assert(model.transform(nearMiss).schema == model.transformSchema(nearMiss.schema))
+        assert(model.transform(nearMiss).collect().length == 2)
+        val renamed = examples.withColumnRenamed("value", "Prediction")
+        val custom = model.copy(org.apache.spark.ml.param.ParamMap.empty).setInputCol("Prediction")
+        if (caseSensitive) {
+          assert(custom.transform(renamed).schema == custom.transformSchema(renamed.schema))
+          assert(custom.transform(renamed).collect().length == 2)
+        } else {
+          intercept[IllegalArgumentException] { custom.transformSchema(renamed.schema) }
+          intercept[IllegalArgumentException] { custom.transform(renamed) }
+        }
+      }
+    } finally {
+      spark.conf.set("spark.sql.caseSensitive", original)
+    }
+  }
+
   private def loadDSJSON = {
     val fileLocation = FileUtilities.join(BuildInfo.datasetDir,
       "VowpalWabbit", "Train", "dsjson_cb_part1.json").toString
@@ -254,6 +383,8 @@ class VerifyVowpalWabbitGeneric extends Benchmarks with EstimatorFuzzing[VowpalW
     val dataset = Seq("1 |a b c", "0 |d e f").toDF("value")
     Seq(new TestObject(
       new VowpalWabbitGeneric(),
-      dataset))
+      dataset),
+      new TestObject(new VowpalWabbitGeneric(),
+        dataset.withColumn("id", F.length($"value")).withColumn("note", F.lit("retained"))))
   }
 }