import sys
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.utils import getResolvedOptions
from pyspark.sql import functions as F
from pyspark.sql.types import DecimalType
from decimal import Decimal, ROUND_FLOOR, ROUND_HALF_UP, ROUND_CEILING


def round_fraction(target: Decimal, meth: str, pos: str):
    p = {
        "1": "1.",
        "2": "0.1",
        "3": "0.01",
        "4": "0.001",
    }
    methods = {
        "1": ROUND_FLOOR,
        "2": ROUND_HALF_UP,
        "3": ROUND_CEILING,
    }
    if meth is None or pos is None or meth == "" or pos == "":
        return target
    return target.quantize(Decimal(p[pos]), rounding=methods[meth])


udf_round_fraction = F.udf(round_fraction, DecimalType(10, 3))


def exec_sample(glueContext, spark, input_dir, output_dir):
    data = ["calc_source", "attributes"]
    for d in data:
        p = f"s3://{input_dir}/{d}/"
        # Sample to convert DynamicFrame to DataFrame then create temporary view
        glueContext.create_dynamic_frame.from_options(
            connection_type="s3",
            connection_options={"paths": [p]},
            format="parquet",
        ).toDF().createOrReplaceTempView(d)

    # SQL sample to join two tables
    wk_main = spark.sql("""
        SELECT
            src.id
        ,	src.number1
        ,	src.number2
        ,	src.position
        ,	src.method
        ,	src.group
        ,	att.attribute1
        FROM
            calc_source	src
        INNER JOIN
            attributes	att
        ON
            src.group	=	att.group
        WHERE
            src.number1	IS	NOT NULL
        AND	src.number2	IS	NOT NULL
        AND	src.number2	<>	0
    """)

    # UDF sample to calculate fractions
    wk_main = wk_main.withColumn(
        "calc_result",
        udf_round_fraction(
            F.col("number1") / F.col("number2"),
            F.col("method"),
            F.col("position"))
    )

    wk_main.write.mode("overwrite").format("parquet").save(f"s3://{output_dir}/sample_out/")


def main():
    args = getResolvedOptions(sys.argv, ["input_dir", "output_dir"])
    glueContext = GlueContext(SparkContext.getOrCreate())
    spark = glueContext.spark_session
    # Exec
    exec_sample(glueContext, spark, args["input_dir"], args["output_dir"])


if __name__ == '__main__':
    main()

from pyspark.context import SparkContext
from awsglue.context import GlueContext


glueContext = GlueContext(SparkContext.getOrCreate())
df = glueContext.create_dynamic_frame.from_catalog(
    database="sampledb",
    table_name="departuredelays",
    push_down_predicate="(any=='sample')",
).toDF()

df.write.mode("overwrite").format("parquet").save(
    "s3://bucket_name/departuredelays_out/any=sample/"
)