from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("Export-Sample-For-Sklearn") \
    .master("spark://10.0.0.15:7077") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")

df_full = spark.read.parquet("hdfs://10.0.0.15:9000/ids/dataset_full_clean_v2.parquet")
df = df_full.sample(withReplacement=False, fraction=0.25, seed=42)
df = df.drop("Timestamp")  # not needed for sklearn either

print("Sample rows: " + str(df.count()))

# Write as a single CSV file to HDFS first, then we'll pull it to local disk
df.coalesce(1).write.mode("overwrite").option("header", "true").csv(
    "hdfs://10.0.0.15:9000/ids/sample_25pct_for_sklearn_csv")

print("Export done.")
spark.stop()
