#!/usr/bin/env python
# coding: utf-8

# >### 🚩 *Create a free WhyLabs account to get more value out of whylogs!*<br> 
# >*Did you know you can store, visualize, and monitor whylogs profiles with the [WhyLabs Observability Platform](https://whylabs.ai/whylogs-free-signup?utm_source=whylogs-Github&utm_medium=whylogs-example&utm_campaign=Pyspark_Profiling)? Sign up for a [free WhyLabs account](https://whylabs.ai/whylogs-free-signup?utm_source=whylogs-Github&utm_medium=whylogs-example&utm_campaign=Pyspark_Profiling) to leverage the power of whylogs and WhyLabs together!*

# # PySpark Integration
# [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/whylabs/whylogs/blob/mainline/python/examples/integrations/Pyspark_Profiling.ipynb)
# 
# 
# Hi! Perhaps you're already feeling confident with our library, but you really wish there was an easy way to plug our profiling into your existing PySpark jobs. Well, glad you've made it here, because this is what we are going to cover in this example notebook 😃
# 
# If you wish to have other insights on how to use whylogs, feel free to check our [other existing examples](https://github.com/whylabs/whylogs/tree/mainline/python/examples), as they might be extremely useful!

# ## Installing the extra dependency
# 
# As we want to enable users to have exactly what they need to use from whylogs, the `pyspark` integration comes as an extra dependency. In order to have it available, simply uncomment and run the following cell:

# In[2]:


# Note: you may need to restart the kernel to use updated packages.
get_ipython().run_line_magic('pip', "install 'whylogs[spark]'")


# ## Initializing a SparkSession
# 
# Here we will initialize a SparkSession. I'm also setting the `pyarrow` execution config, because it makes our methods even more performant. 
# 
# >**IMPORTANT**: Make sure you have Spark 3.0+ available in your environment, as our implementation relies on it for a smoother integration

# In[3]:


from pyspark.sql import SparkSession


# In[ ]:


spark = SparkSession.builder.appName('whylogs-testing').getOrCreate()
arrow_config_key = "spark.sql.execution.arrow.pyspark.enabled"
spark.conf.set(arrow_config_key, "true")


# ## Reading the data
# 
# For the sake of simplicity (and computational efforts, so you can run this notebook from your local machine), we will read the Wine Quality dataset, available in this URL: "http://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv".

# In[5]:


from pyspark import SparkFiles

data_url = "http://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
spark.sparkContext.addFile(data_url)


# In[6]:


spark_dataframe = spark.read.option("delimiter", ";").option("inferSchema", "true").csv(SparkFiles.get("winequality-red.csv"), header=True)


# In[7]:


spark_dataframe.show(n=1, vertical=True)


# In[8]:


spark_dataframe.printSchema()


# ## Profiling the data with whylogs
# 
# Now that we have a Spark DataFrame in place, let's see how easy it is to profile our data with whylogs.

# In[9]:


from whylogs.api.pyspark.experimental import collect_column_profile_views

column_views_dict = collect_column_profile_views(spark_dataframe)


# Yeap. It's done. It is **that** easy.
# 
# But what do we get with a `column_views_dict`? 

# In[10]:


print(column_views_dict)


# It is a dictionary with one `ColumnProfileView` object per column in your dataset. And we can inspect some of the metrics on each one of them, such as the counts for a given column

# In[11]:


column_views_dict["density"].get_metric("counts").n.value, spark_dataframe.count()


# Or their `mean` value:

# In[12]:


column_views_dict["density"].get_metric("distribution").mean.value


# And now let's check how accurate whylogs did store that `mean` calculation.

# In[13]:


from pyspark.sql.functions import mean
spark_dataframe.select(mean("density")).show()


# It is not the literal exact value, but it gets really close, right? That is because we are not extracting the exact information, but we are also **not sampling** the data. `whylogs` will look at **every data point** and *statistically* decide wether or not that data point is relevant to the final calculation. 
# 
# Is it just me or this is extremely powerful? Yes, it is.
# 
# > "Cool! But what can I do with a bunch of `ColumnProfileView`'s from my Dataset? I want to see everything together
# 
# Well, you've come to the right place, because we will inspect the next method that does just that!

# In[14]:


from whylogs.api.pyspark.experimental import collect_dataset_profile_view

dataset_profile_view = collect_dataset_profile_view(input_df=spark_dataframe)


# Yes, that easy. You now have a `DatasetProfileView`. As you might have seen from other example notebooks in our repo, you can turn this *lightweight* object into a pandas DataFrame, and visualize all the important metrics that we've profiled, like this:

# In[15]:


import pandas as pd 

dataset_profile_view.to_pandas().head()


# ## Persisting as a file
# 
# After collecting profiles, it is a good practice to store them as a file. This will allow you to later on read them back, merge with future profiles and track how is your data behaving along the way.

# In[17]:


dataset_profile_view.write(path="my_super_awesome_profile.bin")


# And that's it, you have just written a profile generated with spark to your local environment! If you wish to upload to different locations, such as s3, whylabs or others, please make sure to check out our [other examples](https://github.com/whylabs/whylogs/tree/mainline/python/examples) page.
# 
# Hopefully this tutorial will help you get started to profile and observe your data behaviour in your Spark jobs with almost no friction :)

# ## Important note
# 
# As you might have seen from the imports, currently this pyspark implementation is the **experimental** phase. We ran some benchmark ourselves with it, and for the sake of example, a `90Gb` dataset with 80M rows could be profiled in under 3 minutes! Cool, right? But we still want more users to try this on their own, see if there are places to be improved and give us feedback before we make it officially **the** spark module here. 
# Please, feel free to reach out to our [community Slack](https://communityinviter.com/apps/whylabs-community/rsqrd-ai-community) and interact with us there. We will love to hear from you :)