Fix after review

rules/S7196/python/metadata.json

@@ -1,5 +1,5 @@
 {
-  "title": "Complex logic provided to PySpark withColumn method should be refactored into a separate expression",
+  "title": "Complex logic provided to PySpark \"withColumn\" method should be refactored into a separate expression",
   "type": "CODE_SMELL",
   "status": "ready",
   "remediation": {

rules/S7196/python/rule.adoc

@@ -1,12 +1,12 @@
-This rule raises an issue when complex functions or expressions are directly passed to withColumn
+This rule raises an issue when complex expressions are directly passed to `withColumn`, `filter` or `when` functions.
 
 == Why is this an issue?
 
-`withColumn` method is commonly used to add or modify columns in a DataFrame. When complex functions or expressions are directly passed to withColumn, it can lead to code that is difficult to read, understand, and maintain. Also, it will become easier to write unit tests for these functions, ensuring that the logic is correct and behaves as expected. This leads to more robust and reliable code. 
+`withColumn`, `filter` and `when` methods are commonly used to add, modify, or filter columns in a DataFrame. When long or complex expressions are directly passed to those functions, it can lead to code that is difficult to read, understand, and maintain. Refactoring such expressions into functions or variables will help with readability. Also, it will become easier to write unit tests for these functions, ensuring that the logic is correct and behaves as expected. This leads to more robust and reliable code. 
 
 == How to fix it
 
-To fix this issue, complex logic within `withColumn` logic should be refactored into separate functions or variables before being passed to `withColumn` to improve code clarity and maintainability,
+To fix this issue, complex logic within `withColumn`, `filter` and `when` calls should be refactored into separate functions or variables to improve code clarity and maintainability,
 
 === Code examples
 
@@ -15,28 +15,39 @@ To fix this issue, complex logic within `withColumn` logic should be refactored
 [source,python,diff-id=1,diff-type=noncompliant]
 ----
 from pyspark.sql.functions import *
-df = df.withColumn('Revenue', col('fare_amount').substr(0, 10).cast("float") + col('extra').substr(0, 5).cast("float") + col('tax').substr(0, 3).cast("float"))
+df = df.withColumn('Revenue', col('fare_amount').substr(0, 10).cast("float") + col('extra').substr(0, 5).cast("float") + col('tax').substr(0, 3).cast("float")) # Noncompliant
+df = df.withColumn('High revenue', when(col('fare_amount').substr(0, 10).cast("float") > 100. and col('extra').substr(0, 5).cast("float") > 100. and col('tax').substr(0, 3).cast("float") < 50.)) # Noncompliant
+df = df.filter(col('fare_amount').substr(0, 10).cast("float") > 100. and col('extra').substr(0, 5).cast("float") > 100. and col('tax').substr(0, 3).cast("float") < 50.) # Noncompliant
 ----
 
 ==== Compliant solution
 
 [source,python,diff-id=1,diff-type=compliant]
 ----
 from pyspark.sql.functions import *
-def convert_to_float(col_str):
-    return col_str.substr(0, 10).cast("float")
-    
-def compute_revenue():
-        fare_amount = col('fare_amount').substr(0, 10).cast("float")
+
+def get_revenue_inputs():
+        fare_amount = col('fare_amount').substr(0, 15).cast("float")
         extra = col('extra').substr(0, 5).cast("float")
         tax = col('tax').substr(0, 3).cast("float")
-        
+        return fare_amount, extra, tax 
+
+def compute_revenue():
+        fare_amount, extra, tax = get_revenue_inputs()
         return fare_amount + extra + tax
+
+def is_high_revenue():
+    fare_amount, extra, tax = get_revenue_inputs()
+    return when( (fare_amount > 100.) & (extra > 100.) & (tax < 50), True).otherwise(False)
         
 df = df.withColumn("Revenue", compute_revenue()) # Compliant
+df = df.withColumn('High revenue', is_high_revenue()) # Compliant
+df = df.filter( is_high_revenue() ) # Compliant
 ----
 
 == Resources
 === Documentation
 
- * PySpark withColumn Documentation - https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.withColumn.html[pyspark.sql.DataFrame.withColumn]
+ * PySpark withColumn documentation - https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.withColumn.html[pyspark.sql.DataFrame.withColumn]
+ * PySpark filter documentation - https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.filter.html[pyspark.sql.DataFrame.filter]
+ * PySpark when documentation - https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.when.html[pyspark.sql.functions.when]