<jupyter_start><jupyter_text>VTL Simple analytic function<jupyter_code>from pyspark.sql import SparkSession,DataFrame
from pyspark.sql.types import StructField, StructType, StringType, DoubleType, IntegerType, LongType, DecimalType
import os
from pyspark.sql.functions import lit, count,sum,avg,collect_list,min,max,percentile_approx,stddev_pop,stddev_samp,var_pop,var_samp
from pyspark.sql.window import Window
local = True

if local:
    spark = SparkSession.builder \
        .master("local[4]") \
        .appName("VTLAnalytic")\
        .getOrCreate()
else:
    spark = SparkSession.builder\
        .master("k8s://https://kubernetes.default.svc:443") \
        .appName("VTLAnalytic")\
        .config("spark.kubernetes.container.image", "inseefrlab/jupyter-datascience:py3.9.7-spark3.2.0")\
        .config("spark.kubernetes.authenticate.driver.serviceAccountName", os.environ['KUBERNETES_SERVICE_ACCOUNT'])\
        .config("spark.executor.instances", "4")\
        .config("spark.executor.memory", "8g")\
        .config("spark.kubernetes.namespace", os.environ['KUBERNETES_NAMESPACE'])\
        .getOrCreate()
data=[("A", "XX", 2000, 3, 1),
    ("A", "XX", 2001, 4, 9),
    ("A", "XX", 2002, 7, 5),
    ("A", "XX", 2003, 6, 8),
    ("A", "YY", 2000, 9, 3),
    ("A", "YY", 2001, 5, 4),
    ("A", "YY", 2002, 10, 2),
    ("A", "YY", 2003, 5, 7)]

schema=StructType([StructField("Id_1",StringType(),True),
                   StructField("Id_2",StringType(),True),
                   StructField("Id_3",IntegerType(),True),
                   StructField("Me_1",IntegerType(),True),
                   StructField("Me_2",IntegerType(),True)])

df=spark.createDataFrame(data, schema)
df.show()<jupyter_output><jupyter_text>1.1 Countcount ( DS_1 over ( partition by Id_1 ) )<jupyter_code>partition_col_name="Id_1"
win_name=Window.partitionBy(partition_col_name)
target_col_name="Me_1"
new_col_name=f"count_{target_col_name}"
df_count=df.withColumn(new_col_name,count(target_col_name).over(win_name))
df_count.show()
df.printSchema()<jupyter_output>root
 |-- Id_1: string (nullable = true)
 |-- Id_2: string (nullable = true)
 |-- Id_3: integer (nullable = true)
 |-- Me_1: integer (nullable = true)
 |-- Me_2: integer (nullable = true)<jupyter_text>1.2 SumWe can do a count without order byDS_r :=count ( DS_1 over ( partition by Id_1 ) )<jupyter_code>partition_col_name="Id_1"
win_name=Window.partitionBy(partition_col_name)
target_col_name="Me_1"
new_col_name=f"sum_{target_col_name}"
df_sum=df.withColumn(new_col_name,sum(target_col_name).over(win_name))
df_sum.show()<jupyter_output>+----+----+----+----+----+--------+
|Id_1|Id_2|Id_3|Me_1|Me_2|sum_Me_1|
+----+----+----+----+----+--------+
|   A|  XX|2000|   3|   1|      49|
|   A|  XX|2001|   4|   9|      49|
|   A|  XX|2002|   7|   5|      49|
|   A|  XX|2003|   6|   8|      49|
|   A|  YY|2000|   9|   3|      49|
|   A|  YY|2001|   5|   4|      49|
|   A|  YY|2002|  10|   2|      49|
|   A|  YY|2003|   5|   7|      49|
+----+----+----+----+----+--------+<jupyter_text>We can also do without partitionDS_r := sum ( DS_1 over ( order by Id_1, Id_2, Id_3 ) )<jupyter_code>partition_col_names=["Id_1","Id_2","Id_3"]
win_name=Window.orderBy(partition_col_names)
target_col_name="Me_1"
new_col_name=f"sum_{target_col_name}"
df_sum=df.withColumn(new_col_name,sum(target_col_name).over(win_name))
df_sum.show()
partition_col_name="Id_1"
win_name=Window.partitionBy(partition_col_name).rowsBetween(-1,1)
target_col_name="Me_1"
new_col_name=f"sum_{target_col_name}"
df_collect=df.withColumn(new_col_name,collect_list(target_col_name).over(win_name))
df_collect.show()
partition_col_name=["Id_1","Id_2"]
order_col_name=["Id_3"]
win_name=Window.partitionBy(partition_col_name).orderBy(order_col_name).rangeBetween(-3,3)
target_col_name="Me_1"
new_col_name=f"sum_{target_col_name}"
df_collect1=df.withColumn(new_col_name,collect_list(target_col_name).over(win_name))
df_collect1.show()<jupyter_output>+----+----+----+----+----+-------------+
|Id_1|Id_2|Id_3|Me_1|Me_2|     sum_Me_1|
+----+----+----+----+----+-------------+
|   A|  YY|2000|   9|   3|[9, 5, 10, 5]|
|   A|  YY|2001|   5|   4|[9, 5, 10, 5]|
|   A|  YY|2002|  10|   2|[9, 5, 10, 5]|
|   A|  YY|2003|   5|   7|[9, 5, 10, 5]|
|   A|  XX|2000|   3|   1| [3, 4, 7, 6]|
|   A|  XX|2001|   4|   9| [3, 4, 7, 6]|
|   A|  XX|2002|   7|   5| [3, 4, 7, 6]|
|   A|  XX|2003|   6|   8| [3, 4, 7, 6]|
+----+----+----+----+----+-------------+<jupyter_text>1.3 minDS_r := min ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name).orderBy(order_by_col)
target_col_name="Me_1"
new_col_name=f"min_{target_col_name}"
df_min=df.withColumn(new_col_name,min(target_col_name).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+--------+
|Id_1|Id_2|Id_3|Me_1|Me_2|min_Me_1|
+----+----+----+----+----+--------+
|   A|  YY|2000|   9|   3|       9|
|   A|  YY|2001|   5|   4|       5|
|   A|  YY|2002|  10|   2|       5|
|   A|  YY|2003|   5|   7|       5|
|   A|  XX|2000|   3|   1|       3|
|   A|  XX|2001|   4|   9|       3|
|   A|  XX|2002|   7|   5|       3|
|   A|  XX|2003|   6|   8|       3|
+----+----+----+----+----+--------+<jupyter_text>1.4 maxDS_r := max ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name).orderBy(order_by_col)
target_col_name="Me_1"
new_col_name=f"max_{target_col_name}"
df_min=df.withColumn(new_col_name,max(target_col_name).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+--------+
|Id_1|Id_2|Id_3|Me_1|Me_2|max_Me_1|
+----+----+----+----+----+--------+
|   A|  YY|2000|   9|   3|       9|
|   A|  YY|2001|   5|   4|       9|
|   A|  YY|2002|  10|   2|      10|
|   A|  YY|2003|   5|   7|      10|
|   A|  XX|2000|   3|   1|       3|
|   A|  XX|2001|   4|   9|       4|
|   A|  XX|2002|   7|   5|       7|
|   A|  XX|2003|   6|   8|       7|
+----+----+----+----+----+--------+<jupyter_text>1.5 avgDS_r := avg ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name).orderBy(order_by_col)
target_col_name="Me_1"
new_col_name=f"avg_{target_col_name}"
df_min=df.withColumn(new_col_name,avg(target_col_name).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+-----------------+
|Id_1|Id_2|Id_3|Me_1|Me_2|         avg_Me_1|
+----+----+----+----+----+-----------------+
|   A|  YY|2000|   9|   3|              9.0|
|   A|  YY|2001|   5|   4|              7.0|
|   A|  YY|2002|  10|   2|              8.0|
|   A|  YY|2003|   5|   7|             7.25|
|   A|  XX|2000|   3|   1|              3.0|
|   A|  XX|2001|   4|   9|              3.5|
|   A|  XX|2002|   7|   5|4.666666666666667|
|   A|  XX|2003|   6|   8|              5.0|
+----+----+----+----+----+-----------------+<jupyter_text>1.6 medianDS_r := median ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )percentile_approx with partitionBy followed by orderBy will do rolling median. To have median of each partition, you must only have partitionBy in your window definition.<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name)
target_col_name="Me_1"
new_col_name=f"avg_{target_col_name}"
df_min=df.withColumn(new_col_name,percentile_approx(target_col_name,0.5,10000000).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+--------+
|Id_1|Id_2|Id_3|Me_1|Me_2|avg_Me_1|
+----+----+----+----+----+--------+
|   A|  YY|2000|   9|   3|       5|
|   A|  YY|2001|   5|   4|       5|
|   A|  YY|2002|  10|   2|       5|
|   A|  YY|2003|   5|   7|       5|
|   A|  XX|2000|   3|   1|       4|
|   A|  XX|2001|   4|   9|       4|
|   A|  XX|2002|   7|   5|       4|
|   A|  XX|2003|   6|   8|       4|
+----+----+----+----+----+--------+<jupyter_text>1.7 stddev_popThe operator returns the “population standard deviation” of the input values.DS_r := stddev_pop ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name)
target_col_name="Me_1"
new_col_name=f"stddev_pop_{target_col_name}"
df_min=df.withColumn(new_col_name,stddev_pop(target_col_name).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+------------------+
|Id_1|Id_2|Id_3|Me_1|Me_2|   stddev_pop_Me_1|
+----+----+----+----+----+------------------+
|   A|  YY|2000|   9|   3| 2.277608394786075|
|   A|  YY|2001|   5|   4| 2.277608394786075|
|   A|  YY|2002|  10|   2| 2.277608394786075|
|   A|  YY|2003|   5|   7| 2.277608394786075|
|   A|  XX|2000|   3|   1|1.5811388300841895|
|   A|  XX|2001|   4|   9|1.5811388300841895|
|   A|  XX|2002|   7|   5|1.5811388300841895|
|   A|  XX|2003|   6|   8|1.5811388300841895|
+----+----+----+----+----+------------------+<jupyter_text>1.8 stddev_sampThe operator returns the “sample standard deviation” of the input values.DS_r :=  stddev_stamp ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name)
target_col_name="Me_1"
new_col_name=f"stddev_samp_{target_col_name}"
df_min=df.withColumn(new_col_name,stddev_samp(target_col_name).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+------------------+
|Id_1|Id_2|Id_3|Me_1|Me_2|  stddev_samp_Me_1|
+----+----+----+----+----+------------------+
|   A|  YY|2000|   9|   3|2.6299556396765835|
|   A|  YY|2001|   5|   4|2.6299556396765835|
|   A|  YY|2002|  10|   2|2.6299556396765835|
|   A|  YY|2003|   5|   7|2.6299556396765835|
|   A|  XX|2000|   3|   1|1.8257418583505536|
|   A|  XX|2001|   4|   9|1.8257418583505536|
|   A|  XX|2002|   7|   5|1.8257418583505536|
|   A|  XX|2003|   6|   8|1.8257418583505536|
+----+----+----+----+----+------------------+<jupyter_text>1.9 var_popDS_r :=  var_pop ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )The operator returns the “population variance” of the input values<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name)
target_col_name="Me_1"
new_col_name=f"var_pop_{target_col_name}"
df_min=df.withColumn(new_col_name,var_pop(target_col_name).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+------------------+
|Id_1|Id_2|Id_3|Me_1|Me_2|      var_pop_Me_1|
+----+----+----+----+----+------------------+
|   A|  YY|2000|   9|   3|            5.1875|
|   A|  YY|2001|   5|   4|            5.1875|
|   A|  YY|2002|  10|   2|            5.1875|
|   A|  YY|2003|   5|   7|            5.1875|
|   A|  XX|2000|   3|   1|2.4999999999999996|
|   A|  XX|2001|   4|   9|2.4999999999999996|
|   A|  XX|2002|   7|   5|2.4999999999999996|
|   A|  XX|2003|   6|   8|2.4999999999999996|
+----+----+----+----+----+------------------+<jupyter_text>1.10 var_sampDS_r :=  var_samp ( DS_1 over ( partition by Id_1, Id_2 order by Id_3 ) )The operator returns the “sample variance” of the input values<jupyter_code>partition_col_names=["Id_1","Id_2"]
order_by_col=["Id_3"]
win_name=Window.partitionBy(partition_col_name)
target_col_name="Me_1"
new_col_name=f"var_samp_{target_col_name}"
df_min=df.withColumn(new_col_name,var_samp(target_col_name).over(win_name))
df_min.show()<jupyter_output>+----+----+----+----+----+------------------+
|Id_1|Id_2|Id_3|Me_1|Me_2|     var_samp_Me_1|
+----+----+----+----+----+------------------+
|   A|  YY|2000|   9|   3| 6.916666666666667|
|   A|  YY|2001|   5|   4| 6.916666666666667|
|   A|  YY|2002|  10|   2| 6.916666666666667|
|   A|  YY|2003|   5|   7| 6.916666666666667|
|   A|  XX|2000|   3|   1|3.3333333333333326|
|   A|  XX|2001|   4|   9|3.3333333333333326|
|   A|  XX|2002|   7|   5|3.3333333333333326|
|   A|  XX|2003|   6|   8|3.3333333333333326|
+----+----+----+----+----+------------------+
{'fraction_non_alphanumeric': 0.24097980801059252, 'fraction_numerical': 0.11105594174114532, 'mean_word_length': 2.7127496159754223, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 14, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '501732', 'license': 'MIT', 'n_tokens_mistral': 5982, 'n_tokens_neox': 5278, 'n_words': 816, 'path': 'notebooks/VTLAnalyticSimple.ipynb', 'repo_name': 'pengfei99/VTLPerformanceTest', 'size': '12084'}