Tutorial: Get started with Snowpark Connect for Spark¶

import com.snowflake.snowpark_connect.client.SnowparkConnectSession;
import com.snowflake.snowpark_connect.client.SnowflakeSession;
import org.apache.spark.sql.*;
import org.apache.spark.sql.types.*;
import org.apache.spark.sql.connect.client.REPLClassDirMonitor;
import static org.apache.spark.sql.functions.*;

public class Tutorial {
    public static void main(String[] args) {
        SparkSession spark = SnowparkConnectSession.builder()
            .appName("Tutorial")
            .getOrCreate();

        // Register a class finder so Spark can discover compiled UDF classes
        REPLClassDirMonitor classFinder = new REPLClassDirMonitor("/path/to/target/classes");
        spark.registerClassFinder(classFinder);

        spark.sql("CREATE OR REPLACE TABLE sales (product STRING, region STRING, amount DOUBLE)");
        spark.sql("INSERT INTO sales VALUES "
            + "('Widget', 'north', 120.50), "
            + "('Gadget', 'south', 85.00), "
            + "('Widget', 'south', 200.75), "
            + "('Gizmo', 'north', 45.30), "
            + "('Gadget', 'north', 150.00), "
            + "('Gizmo', 'south', 60.00)");

import org.apache.spark.api.java.function.MapFunction;

// POJO that mirrors the source table schema.
// Encoders.bean() maps DataFrame columns to getter/setter pairs,
// giving you a typed Dataset<Sale> instead of raw Row objects.
public static class Sale implements java.io.Serializable {
    private String product;
    private String region;
    private double amount;

    public Sale() {}
    public String getProduct() { return product; }
    public void setProduct(String p) { this.product = p; }
    public String getRegion() { return region; }
    public void setRegion(String r) { this.region = r; }
    public double getAmount() { return amount; }
    public void setAmount(double a) { this.amount = a; }
}

// Output POJO with the computed column. A separate class lets the
// encoder infer the result schema automatically from the bean fields.
public static class SaleWithTax implements java.io.Serializable {
    private String product;
    private String region;
    private double amount;
    private double amount_with_tax;

    public SaleWithTax() {}
    public String getProduct() { return product; }
    public void setProduct(String p) { this.product = p; }
    public String getRegion() { return region; }
    public void setRegion(String r) { this.region = r; }
    public double getAmount() { return amount; }
    public void setAmount(double a) { this.amount = a; }
    public double getAmount_with_tax() { return amount_with_tax; }
    public void setAmount_with_tax(double v) { this.amount_with_tax = v; }
}

static double applyTieredTax(double amount) {
    double rate = amount < 50 ? 0.05 : amount <= 150 ? 0.10 : 0.15;
    return Math.round(amount * (1 + rate) * 100.0) / 100.0;
}

Dataset<Sale> sales = spark.table("sales").as(Encoders.bean(Sale.class));

Dataset<Row> result = sales
    .map((MapFunction<Sale, SaleWithTax>) s -> {
        SaleWithTax out = new SaleWithTax();
        out.setProduct(s.getProduct());
        out.setRegion(s.getRegion().toUpperCase());
        out.setAmount(s.getAmount());
        out.setAmount_with_tax(applyTieredTax(s.getAmount()));
        return out;
    }, Encoders.bean(SaleWithTax.class))
    .toDF();

result.show();

result.write().mode("overwrite").saveAsTable("sales_with_tax");

result.write().mode("overwrite").parquet("@~/tutorial_output/sales_with_tax");

        SnowflakeSession sf = new SnowflakeSession(spark);

        Dataset<Row> totals = sf.sql(
            "SELECT product, "
            + "SUM(amount_with_tax) AS total_with_tax, "
            + "COUNT(*) AS num_sales "
            + "FROM sales_with_tax "
            + "GROUP BY product "
            + "ORDER BY total_with_tax DESC");
        totals.show();

        Dataset<Row> stagedFiles = sf.sql(
            "LIST @~/tutorial_output/sales_with_tax");
        stagedFiles.show(false);

        spark.stop();
    }
}