fix: Update SparkKafkaProcessor to drop unnecessary columns and renam…

…e feature_value to value
ExpediaGroup · Jan 15, 2025 · 95bc2a6 · 95bc2a6
1 parent 075c4c0
commit 95bc2a6
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/sdk/python/feast/infra/contrib/spark_kafka_processor.py b/sdk/python/feast/infra/contrib/spark_kafka_processor.py
@@ -338,9 +338,11 @@ def batch_write_with_connector(
             batch_id: int,
         ):
             start_time = time.time()
-            sdf = sdf.drop("event_header")
             convert_to_blob = udf(lambda s: s.encode("utf-8"), BinaryType())
-            sdf = sdf.withColumn("feature_value", convert_to_blob(col("feature_value")))
+            sdf = sdf.withColumn("value", convert_to_blob(col("feature_value"))).drop(
+                "event_header",
+                "feature_value",
+            )
             sdf.write.format("org.apache.spark.sql.cassandra").mode("append").options(
                 table="mlpfs_scylladb_perf_test_cc_stream_fv", keyspace="feast"
             ).save()