From c954205316e73f1f117626c8a289f0f6d2a3c3b5 Mon Sep 17 00:00:00 2001 From: Lenin Mookiah <2639584+leninworld@users.noreply.github.com> Date: Mon, 9 Nov 2026 02:49:48 -0800 Subject: [PATCH 1/2] Preserve input schemas in VW Generic predictions #2817 Keep all input fields and append the native prediction schema, restoring top-level metadata after row encoding. Reject prediction-name collisions using the Spark session case-sensitivity setting. Cover schema ordering, scalar and array predictions, custom inputs and collision behavior. --- .../synapse/ml/vw/VowpalWabbitGeneric.scala | 54 +++++++++--- .../ml/vw/VerifyVowpalWabbitGeneric.scala | 87 ++++++++++++++++++- 2 files changed, 130 insertions(+), 11 deletions(-) diff --git a/vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala b/vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala index 7537a3511f0..5c606684408 100644 --- a/vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala +++ b/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,9 +101,27 @@ 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 top-level input fields with their type, nullability and metadata, and 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 ({ @@ -105,9 +129,11 @@ class VowpalWabbitGenericModel(override val uid: String) 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 +142,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).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) } } diff --git a/vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala b/vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala index 175b61dcce2..6873a2b3fef 100644 --- a/vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala +++ b/vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala @@ -240,6 +240,89 @@ 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 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 +337,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")))) } } From d24b27872c2b8bf8be0c8d3715b4ec59894e0954 Mon Sep 17 00:00:00 2001 From: Rana Singh Date: Mon, 9 Nov 2026 04:43:29 -0800 Subject: [PATCH 2/2] fix: preserve nested metadata in VW Generic output ## Summary Restore each output column's declared datatype after Spark row encoding so nested field metadata survives VowpalWabbitGenericModel.transform. Add public Spark regressions for nested schemas and categorical label decoding. ## Prompting Intent Apply the independently reproduced nested-metadata correction to the existing external contribution, preserve the contributor's implementation and history, and validate the improvement through the PR and external-contributor loops. ## Linked Sources - Original report: https://github.com/microsoft/SynapseML/issues/2817 - Contributor PR: https://github.com/microsoft/SynapseML/pull/2821 - Review finding: https://github.com/microsoft/SynapseML/pull/2821#discussion_r4511954555 - Spark row encoder: https://github.com/apache/spark/blob/v3.5.0/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala - Spark serializer: https://github.com/apache/spark/blob/v3.5.0/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SerializerBuildHelper.scala - Spark casts: https://github.com/apache/spark/blob/v3.5.0/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala - Label decoder: https://github.com/apache/spark/blob/v3.5.0/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala ## Rationale Top-level aliases cannot recover metadata on nested struct fields erased by ExpressionEncoder. Casting to the original datatype restores that schema while retaining the existing top-level metadata alias. Spark treats structurally identical casts as representation-preserving, avoiding recursive reconstruction, new native calls, or shuffles. Public tests check structs, arrays, map keys and values, nullability, collected values, and downstream IndexToString behavior. No public signatures, parameter formats, dependency pins, or CI files change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ed340ce9-385a-41e0-86c4-fe0f09097478 --- .../synapse/ml/vw/VowpalWabbitGeneric.scala | 5 +- .../ml/vw/VerifyVowpalWabbitGeneric.scala | 46 +++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala b/vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala index 5c606684408..98638244bbd 100644 --- a/vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala +++ b/vw/src/main/scala/com/microsoft/azure/synapse/ml/vw/VowpalWabbitGeneric.scala @@ -118,7 +118,8 @@ class VowpalWabbitGenericModel(override val uid: String) } /** - * Keeps all top-level input fields with their type, nullability and metadata, and appends flat predictions. + * 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. */ @@ -145,7 +146,7 @@ class VowpalWabbitGenericModel(override val uid: String) 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).as(field.name, field.metadata) + case (field, internalName) => F.col(internalName).cast(field.dataType).as(field.name, field.metadata) }: _*) }, dataset.columns.length) } diff --git a/vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala b/vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala index 6873a2b3fef..347a2c78944 100644 --- a/vw/src/test/scala/com/microsoft/azure/synapse/ml/vw/VerifyVowpalWabbitGeneric.scala +++ b/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} @@ -286,6 +288,50 @@ class VerifyVowpalWabbitGeneric extends Benchmarks with EstimatorFuzzing[VowpalW } } + 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)