diff --git a/examples/pyspark_native/possible_api.py b/examples/pyspark_native/possible_api.py new file mode 100644 index 00000000..16dc67ef --- /dev/null +++ b/examples/pyspark_native/possible_api.py @@ -0,0 +1,131 @@ +import pyspark +from pyspark.sql import SparkSession +from pyspark.sql import functions as F + +from hamilton.function_modifiers import extract_columns, tag + + +# Data Loading +# Filtering is part of data loading -- do we also expose columns like this? +@extract_columns( + *["l_quantity", "l_extendedprice", "l_discount", "l_tax", "l_returnflag", "l_linestatus"] +) +def lineitem( + sc: SparkSession, + path: str, + filter: str = "l_shipdate <= date '1998-12-01' - interval '90' day", +) -> pyspark.sql.DataFrame: + """Loads and filters data from the lineitem table""" + ds: pyspark.sql.DataFrame = ( + sc.read.option("inferSchema", True).option("header", True).csv(path, sep="|") + ) + if filter: + ds = ds.filter(filter) + print(ds.schema) + return ds + + +# transforms we want people to write +def disc_price(l_extendedprice: float, l_discount: float) -> float: + """Computes the discounted price""" + return l_extendedprice * (1 - l_discount) + + +def charge(l_extendedprice: float, l_discount: float, l_tax: float) -> float: + """Computes the charge""" + return l_extendedprice * (1 - l_discount) * (1 + l_tax) + + +# hacking things in via tags +@tag(group_by="l_returnflag,l_linestatus") +def grouped_lineitem( + l_quantity: pyspark.sql.Column, + l_extendedprice: pyspark.sql.Column, + disc_price: pyspark.sql.Column, # do we do some optional syntax here? + # and at run time we check if the column exists, and if so use it. Else skip it. + # basically it means that someone could write one function and determine if all are required or not. + # and save them from having to update this function if they don't want a particular column being passed + # through -- thought downstream of this, they would have to deal with it... so maybe not that valuable? + charge: pyspark.sql.Column, + l_discount: pyspark.sql.Column, + l_returnflag: pyspark.sql.Column, + l_linestatus: pyspark.sql.Column, +) -> pyspark.sql.GroupedData: + """This function declares the "schema" the datastream needs to have via it's arguments. + The body is blank because the graph adapter knows the actual logic to perform. + + Alternate syntax could be to have the decorator declare what's required, the function then takes the datastream, the + group by happens in the function... + """ + pass + + +# hack to get around https://github.com/marsupialtail/quokka/issues/23 +# @tag(materialize="True") +def compute_aggregates(grouped_lineitem: pyspark.sql.GroupedData) -> pyspark.sql.DataFrame: + # Thought: change these into individual functions? + agg_map = { + "l_quantity": ["sum", "avg"], + "l_extendedprice": ["sum", "avg"], + "disc_price": ["sum"], + "charge": ["sum"], + "l_discount": ["avg"], + "*": ["count"], + } + agg_args = [] + for column_name, aggregates in agg_map.items(): + for aggregate in aggregates: + func = getattr(F, aggregate) + agg_args.append(func(F.column(column_name))) + df = grouped_lineitem.agg( + # { + # # "l_quantity": ["sum", "avg"], + # "l_quantity": "sum", + # # "l_extendedprice": ["sum", "avg"], + # "l_extendedprice": "avg", + # "disc_price": "sum", + # "charge": "sum", + # "l_discount": "avg", + # "*": "count", + # } + *agg_args + ) + rename_map = { + "l_returnflag": "al_returnflag", + "l_linestatus": "al_linestatus", + "count(1)": "row_count", + "avg(l_quantity)": "l_quantity_mean", + "sum(l_quantity)": "l_quantity_sum", + "avg(l_extendedprice)": "l_extendedprice_mean", + "sum(l_extendedprice)": "l_extendedprice_sum", + "sum(disc_price)": "disc_price_sum", + "sum(charge)": "charge_sum", + "avg(l_discount)": "l_discount_mean", + } + for old_name, new_name in rename_map.items(): + df = df.withColumnRenamed(old_name, new_name) + return df + + +# this doesn't seem like the right thing: +# def l_quantity_sum(grouped_lineitem: GroupedDataStream) -> pyspark.sql.Column: +# pass + + +# hack to get around `@tag_outputs` not working with `@extract_columns` as expected. +@extract_columns( + *[ + "row_count", + "l_quantity_sum", + "l_extendedprice_sum", + "disc_price_sum", + "charge_sum", + "l_quantity_mean", + "l_extendedprice_mean", + "l_discount_mean", + "al_returnflag", + "al_linestatus", + ] +) +def extract_aggregates(compute_aggregates: pyspark.sql.DataFrame) -> pyspark.sql.DataFrame: + return compute_aggregates diff --git a/examples/pyspark_native/pyspark_adapter.py b/examples/pyspark_native/pyspark_adapter.py new file mode 100644 index 00000000..39985395 --- /dev/null +++ b/examples/pyspark_native/pyspark_adapter.py @@ -0,0 +1,115 @@ +import inspect +from typing import Any, Callable, Dict, Tuple, Type + +from pyspark.sql import Column, DataFrame, GroupedData, types +from pyspark.sql.functions import column, udf + +from hamilton import base, node + + +class PySparkGraphAdapter(base.SimplePythonDataFrameGraphAdapter): + def __init__(self, result_builder: base.ResultMixin = base.DictResult()): + self.df_objects = {} + self.call_count = 0 + self.result_builder = result_builder + + @staticmethod + def check_input_type(node_type: Type, input_value: Any) -> bool: + return True + + @staticmethod + def check_node_type_equivalence(node_type: Type, input_type: Type) -> bool: + return True + + def _lambda_udf(self, df: DataFrame, hamilton_udf: Callable) -> DataFrame: + sig = inspect.signature(hamilton_udf) + input_parameters = dict(sig.parameters) + return_type = sig.return_annotation + print("lambda inputs", input_parameters, return_type, hamilton_udf.__name__) + if return_type == float: + spark_return_type = types.DoubleType() + else: + raise ValueError(f"Unsupported return type {return_type}") + spark_udf = udf(hamilton_udf, spark_return_type) + return df.withColumn( + hamilton_udf.__name__, spark_udf(*[column(name) for name in sig.parameters.keys()]) + ) + + def _sanitize_kwargs(self, kwargs: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Sanitizes the kwargs to remove the datastream node name.""" + df_names = {} + actual_kwargs = {} + for kwarg_key, kwarg_value in kwargs.items(): + if isinstance(kwarg_value, dict) and "__df_name__" in kwarg_value: + df_names[kwarg_key] = kwarg_value["__df_name__"] + actual_kwargs[kwarg_key] = kwarg_value["__result__"] + else: + actual_kwargs[kwarg_key] = kwarg_value + return actual_kwargs, df_names + + def execute_node(self, node: node.Node, kwargs: Dict[str, Any]) -> Any: + self.call_count += 1 + actual_kwargs, df_names = self._sanitize_kwargs(kwargs) + df_name_set = set(df_names.values()) + if node.type == DataFrame: + # assumption is types into this function are only scalars, or other DataFrame/GroupedData objects + df: DataFrame = node.callable(**actual_kwargs) + self.df_objects[node.name] = df + return {"__df_name__": node.name, "__result__": df} + elif node.type == GroupedData: + print("got group by", self.call_count, node.name) + assert len(df_name_set) == 1, f"Error groupby got multiple DataFrames {df_names}" + df = self.df_objects[df_names.popitem()[1]] + print("before select", df.schema) + df = df.select(list(node.input_types.keys())) + print("after select", df.schema) + group_by_cols = node.tags["group_by"].split(",") + df: GroupedData = df.groupby(group_by_cols) + self.df_objects[node.name] = df + return {"__df_name__": node.name, "__result__": df} + elif ( + node.type == Column + and len(node.input_types) == 1 + and node.tags.get("__generated_by__", None) == "extract_columns" + ): + assert ( + len(df_name_set) == 1 + ), f"Error extract_columns got multiple DataFrames {df_names}" + df_name = df_name_set.pop() + print(node.name, node.tags) + # print('got extract_columns', self.call_count, node.name, kwargs) + # get global one + df = self.df_objects[df_name] + return {"__df_name__": df_name, "__result__": df} + else: + assert len(df_name_set) == 1, f"Error udf got multiple DataFrames {df_names}" + print("got udf", self.call_count, node.name, kwargs) + df_name = df_name_set.pop() + df: DataFrame = self.df_objects[df_name] + print(df.schema) + df = self._lambda_udf(df, node.callable) + self.df_objects[df_name] = df + return {"__df_name__": df_name, "__result__": df} + + def build_result(self, **outputs: Dict[str, Any]) -> DataFrame: + """Builds the result and brings it back to this running process. + + :param outputs: the dictionary of key -> Union[ray object reference | value] + :return: The type of object returned by self.result_builder. + """ + requested_ds_set = set(outputs.keys()) + actual_outputs, df_names = self._sanitize_kwargs(outputs) + df_name_set = set(df_names.values()) + assert ( + len(df_name_set) == 1 + ), f"Error got multiple DataStreams to build result from {df_names}" + df = self.df_objects[df_name_set.pop()] + df = df.select(list(actual_outputs.keys())) + global_ds_set = set(df.columns) + if requested_ds_set.intersection(global_ds_set) != requested_ds_set: + raise ValueError( + f"Error: requested columns not found in final dataframe. " + f"Missing: {requested_ds_set.difference(global_ds_set)}." + ) + print("final schema", df.schema) + return df diff --git a/examples/pyspark_native/requirements.txt b/examples/pyspark_native/requirements.txt new file mode 100644 index 00000000..f5b0776c --- /dev/null +++ b/examples/pyspark_native/requirements.txt @@ -0,0 +1,2 @@ +pyspark +sf-hamilton diff --git a/examples/pyspark_native/run_adapter.py b/examples/pyspark_native/run_adapter.py new file mode 100644 index 00000000..a6902948 --- /dev/null +++ b/examples/pyspark_native/run_adapter.py @@ -0,0 +1,57 @@ +import possible_api +import pyspark_adapter +from pyspark.sql import DataFrame, SparkSession +from pyspark.sql import functions as F +from pyspark.sql import types + +from hamilton import driver, log_setup + + +def disc_price(l_extendedprice, l_discount) -> float: + """Computes the discounted price""" + return l_extendedprice * (1 - l_discount) + + +def grr(): + path = "/Users/stefankrawczyk/Downloads/tpc-h-public/lineitem.tbl" + sc = SparkSession.builder.getOrCreate() + sc.sparkContext.setLogLevel("WARN") + log_setup.setup_logging() + df: DataFrame = sc.read.option("inferSchema", True).option("header", True).csv(path, sep="|") + df = df.filter("l_shipdate <= date '1998-12-01' - interval '90' day") + print(df.schema) + spark_udf = F.udf(disc_price, types.FloatType()) + df.withColumn("disc_price", spark_udf(df["l_extendedprice"], df["l_discount"])).show() + + +def main(): + log_setup.setup_logging(log_level=log_setup.LOG_LEVELS["INFO"]) + spark = SparkSession.builder.getOrCreate() + path = "/Users/stefankrawczyk/Downloads/tpc-h-public/lineitem.tbl" + adapter = pyspark_adapter.PySparkGraphAdapter() + dr = driver.Driver({"sc": spark, "path": path}, possible_api, adapter=adapter) + outputs = [ + "al_returnflag", # can comment one of these out and they are dropped + "al_linestatus", + "row_count", + "charge_sum", + "disc_price_sum", + "l_discount_mean", + "l_extendedprice_mean", + "l_extendedprice_sum", + "l_quantity_mean", + "l_quantity_sum", + # "grouped_lineitem" -- calling collect on groups doesn't work it seems... + # "l_returnflag", # can also get intermediate results + # "charge", + # "disc_price", + ] + dr.visualize_execution(outputs, "./my_dag.dot", {}) + df = dr.execute(outputs) + df.show() + spark.stop() + + +if __name__ == "__main__": + main() + # grr() diff --git a/examples/quokka/nothing_special.py b/examples/quokka/nothing_special.py new file mode 100644 index 00000000..14ef7f8e --- /dev/null +++ b/examples/quokka/nothing_special.py @@ -0,0 +1,195 @@ +"""Shows how you can use Hamilton with Quokka today. + +There isn't any transform re-use due to relying on mutation order +of the central DataStream object. +""" +from pyquokka.datastream import DataStream, GroupedDataStream +from pyquokka.df import QuokkaContext + + +def lineitem(qc: QuokkaContext, path: str) -> DataStream: + """Loads and filters data from the lineitem table""" + ds = qc.read_csv(path + "lineitem.tbl", sep="|", has_header=True) + return ds.filter("l_shipdate <= date '1998-12-01' - interval '90' day") + + +def lineitem_mutation_1(lineitem: DataStream) -> DataStream: + return lineitem.with_column( + "disc_price", + lambda x: x["l_extendedprice"] * (1 - x["l_discount"]), + required_columns={"l_extendedprice", "l_discount"}, + ) + + +def lineitem_mutation_2(lineitem_mutation_1: DataStream) -> DataStream: + return lineitem_mutation_1.with_column( + "charge", + lambda x: x["l_extendedprice"] * (1 - x["l_discount"]) * (1 + x["l_tax"]), + required_columns={"l_extendedprice", "l_discount", "l_tax"}, + ) + + +def grouped_lineitem(lineitem_mutation_2: DataStream) -> GroupedDataStream: + return lineitem_mutation_2.groupby( + ["l_returnflag", "l_linestatus"], orderby=["l_returnflag", "l_linestatus"] + ) + + +def computed_aggregates(grouped_lineitem: GroupedDataStream) -> DataStream: + return grouped_lineitem.agg( + { + "l_quantity": ["sum", "avg"], + "l_extendedprice": ["sum", "avg"], + "disc_price": "sum", + "charge": "sum", + "l_discount": "avg", + "*": "count", + } + ) + + +def manual_hello_world_hamilton(qc: QuokkaContext, path: str): + d = lineitem(qc, path) + d = lineitem_mutation_1(d) + d = lineitem_mutation_2(d) + f = grouped_lineitem(d) + df = computed_aggregates(f) + return df.collect() + + +def hello_world_hamilton(qc: QuokkaContext, path: str): + from hamilton import ad_hoc_utils, base, driver + + temp_module = ad_hoc_utils.create_temporary_module( + lineitem, lineitem_mutation_1, lineitem_mutation_2, grouped_lineitem, computed_aggregates + ) + adapter = base.SimplePythonGraphAdapter(base.DictResult()) + dr = driver.Driver({"path": path, "qc": qc}, temp_module, adapter=adapter) + # dr.visualize_execution(['computed_aggregates'], './my_dag.dot', {}) + result = dr.execute(["computed_aggregates"]) + return result["computed_aggregates"].collect() + + +def hello_world_main(qc, path: str): + lineitem = qc.read_csv(path + "lineitem.tbl", sep="|", has_header=True) + d = lineitem.filter("l_shipdate <= date '1998-12-01' - interval '90' day") + print(len(d.schema)) + d = d.with_column( + "disc_price", + lambda x: x["l_extendedprice"] * (1 - x["l_discount"]), + required_columns={"l_extendedprice", "l_discount"}, + ) + print(len(d.schema)) + d = d.with_column( + "charge", + lambda x: x["l_extendedprice"] * (1 - x["l_discount"]) * (1 + x["l_tax"]), + required_columns={"l_extendedprice", "l_discount", "l_tax"}, + ) + print(len(d.schema)) + f = d.groupby(["l_returnflag", "l_linestatus"], orderby=["l_returnflag", "l_linestatus"]).agg( + { + "l_quantity": ["sum", "avg"], + "l_extendedprice": ["sum", "avg"], + "disc_price": "sum", + "charge": "sum", + "l_discount": "avg", + "*": "count", + } + ) + print(f.schema) + print(len(f.schema)) + df = f.collect() + print(df.columns) + print(len(df.columns)) + return df + + +""" +lineitem = qc.read_csv(disk_path + "lineitem.tbl", sep="|", has_header=True) +orders = qc.read_csv(disk_path + "orders.tbl", sep="|", has_header=True) +customer = qc.read_csv(disk_path + "customer.tbl",sep = "|", has_header=True) +part = qc.read_csv(disk_path + "part.tbl", sep = "|", has_header=True) +supplier = qc.read_csv(disk_path + "supplier.tbl", sep = "|", has_header=True) +partsupp = qc.read_csv(disk_path + "partsupp.tbl", sep = "|", has_header=True) +nation = qc.read_csv(disk_path + "nation.tbl", sep = "|", has_header=True) +region = qc.read_csv(disk_path + "region.tbl", sep = "|", has_header=True) +""" + + +def do_12(qc, path: str): + lineitem = qc.read_csv(path + "lineitem.tbl", sep="|", has_header=True) + orders = qc.read_csv(path + "orders.tbl", sep="|", has_header=True) + d = lineitem.join(orders, left_on="l_orderkey", right_on="o_orderkey") + d = d.filter( + "l_shipmode IN ('MAIL','SHIP') and l_commitdate < l_receiptdate and l_shipdate < l_commitdate and \ + l_receiptdate >= date '1994-01-01' and l_receiptdate < date '1995-01-01'" + ) + d = d.with_column( + "high", + lambda x: (x["o_orderpriority"] == "1-URGENT") | (x["o_orderpriority"] == "2-HIGH"), + required_columns={"o_orderpriority"}, + ) + d = d.with_column( + "low", + lambda x: (x["o_orderpriority"] != "1-URGENT") & (x["o_orderpriority"] != "2-HIGH"), + required_columns={"o_orderpriority"}, + ) + f = d.groupby("l_shipmode").aggregate(aggregations={"high": ["sum"], "low": ["sum"]}) + return f.collect() + + +def do_3(qc, path: str): + lineitem: DataStream = qc.read_csv(path + "lineitem.tbl", sep="|", has_header=True) + lineitem = lineitem.select( + [ + "l_quantity", + "l_extendedprice", + "l_discount", + "l_tax", + "l_returnflag", + "l_linestatus", + "l_orderkey", + "l_shipdate", + ] + ) + orders = qc.read_csv(path + "orders.tbl", sep="|", has_header=True) + orders = orders.select(["o_orderkey", "o_orderdate", "o_shippriority", "o_custkey"]) + customer = qc.read_csv(path + "customer.tbl", sep="|", has_header=True) + customer = customer.select(["c_custkey", "c_mktsegment"]) + d = lineitem.join(orders, left_on="l_orderkey", right_on="o_orderkey") + d = customer.join(d, left_on="c_custkey", right_on="o_custkey") + d = d.select( + [ + "l_orderkey", + "c_mktsegment", + "o_orderdate", + "l_shipdate", + "l_extendedprice", + "l_discount", + "o_shippriority", + ] + ) + # d = d.filter("c_mktsegment = 'BUILDING' and o_orderdate < date '1995-03-15' and l_shipdate > date '1995-03-15'") + d = d.with_column( + "revenue", + lambda x: x["l_extendedprice"] * (1 - x["l_discount"]), + required_columns={"l_extendedprice", "l_discount"}, + ) + d = d.select(["revenue", "o_orderdate", "o_shippriority", "l_orderkey"]) + f = d.groupby(["l_orderkey", "o_orderdate", "o_shippriority"]).agg({"revenue": ["sum"]}) + return f.collect() + + +if __name__ == "__main__": + qc_ = QuokkaContext() + path_ = "/Users/stefankrawczyk/Downloads/tpc-h-public/" + # df = hello_world_main(qc_, path_) + # print(df) + # df = hello_world_hamilton(qc_, path_) + # print(df) + # df = manual_hello_world_hamilton(qc_, path_) + df = do_3(qc_, path_) + print(df) + print(len(df)) + print(len(df.columns)) + print(df.columns) diff --git a/examples/quokka/possible_api.py b/examples/quokka/possible_api.py new file mode 100644 index 00000000..4acdfaf9 --- /dev/null +++ b/examples/quokka/possible_api.py @@ -0,0 +1,274 @@ +import polars as pl +from pyquokka.datastream import DataStream, GroupedDataStream +from pyquokka.df import QuokkaContext + +from hamilton.function_modifiers import extract_columns, tag + + +# Data Loading +# Filtering is part of data loading -- do we also expose columns like this? +@extract_columns( + *[ + "l_quantity", + "l_extendedprice", + "l_discount", + "l_tax", + "l_returnflag", + "l_linestatus", + "l_orderkey", + "l_shipdate", + ] +) +def lineitem_table( + qc: QuokkaContext, + path: str, + lineitem_filter: str = None, +) -> DataStream: + """Loads and filters data from the lineitem table""" + ds: DataStream = qc.read_csv(path + "lineitem.tbl", sep="|", has_header=True) + if lineitem_filter: + ds = ds.filter(lineitem_filter) + print(sorted(ds.schema)) + return ds + + +@extract_columns(*["o_orderdate", "o_orderkey", "o_shippriority", "o_custkey"]) +def orders_table( + qc: QuokkaContext, + path: str, + orders_filter: str = None, +) -> DataStream: + """Loads and filters data from the orders table""" + ds: DataStream = qc.read_csv(path + "orders.tbl", sep="|", has_header=True) + if orders_filter: + ds = ds.filter(orders_filter) + print(sorted(ds.schema)) + return ds + + +@extract_columns( + *[ + "c_custkey", + "c_mktsegment", + ] +) +def customer_table( + qc: QuokkaContext, + path: str, + customer_filter: str = None, +) -> DataStream: + """Loads and filters data from the customer table""" + ds: DataStream = qc.read_csv(path + "customer.tbl", sep="|", has_header=True) + if customer_filter: + ds = ds.filter(customer_filter) + print(sorted(ds.schema)) + return ds + + +# transforms we want people to write +def disc_price(l_extendedprice: pl.Series, l_discount: pl.Series) -> pl.Series: + """Computes the discounted price""" + return l_extendedprice * (1 - l_discount) + + +def charge(l_extendedprice: pl.Series, l_discount: pl.Series, l_tax: pl.Series) -> pl.Series: + """Computes the charge""" + return l_extendedprice * (1 - l_discount) * (1 + l_tax) + + +def high_orderpriority(o_orderpriority: pl.Series) -> pl.Series: + return o_orderpriority == "1-URGENT" | o_orderpriority == "2-HIGH" + + +def low_orderpriority(o_orderpriority: pl.Series) -> pl.Series: + return o_orderpriority != "1-URGENT" & o_orderpriority != "2-HIGH" + + +def clo_revenue(clo_extendedprice: pl.Series, clo_discount: pl.Series) -> pl.Series: + return clo_extendedprice * (1 - clo_discount) + + +# hacking things in via tags +@tag(group_by="l_returnflag,l_linestatus", order_by="l_returnflag,l_linestatus") +def grouped_lineitem( + l_quantity: pl.Series, + l_extendedprice: pl.Series, + disc_price: pl.Series, # do we do some optional syntax here? + # and at run time we check if the column exists, and if so use it. Else skip it. + # basically it means that someone could write one function and determine if all are required or not. + # and save them from having to update this function if they don't want a particular column being passed + # through -- thought downstream of this, they would have to deal with it... so maybe not that valuable? + charge: pl.Series, + l_discount: pl.Series, + l_returnflag: pl.Series, + l_linestatus: pl.Series, +) -> GroupedDataStream: + """This function declares the "schema" the datastream needs to have via it's arguments. + The body is blank because the graph adapter knows the actual logic to perform. + + Alternate syntax could be to have the decorator declare what's required, the function then takes the datastream, the + group by happens in the function... + """ + pass + + +# hack to get around https://github.com/marsupialtail/quokka/issues/23 +# @tag(materialize="True") +@extract_columns( + *[ + "__count_sum", + "l_quantity_sum", + "l_extendedprice_sum", + "disc_price_sum", + "charge_sum", + "l_quantity_mean", + "l_extendedprice_mean", + "l_discount_mean", + "al_returnflag", + "al_linestatus", + ] +) +def compute_aggregates(grouped_lineitem: GroupedDataStream) -> DataStream: + # Thought: change these into individual functions? + return grouped_lineitem.agg( + { + "l_quantity": ["sum", "avg"], + "l_extendedprice": ["sum", "avg"], + "disc_price": "sum", + "charge": "sum", + "l_discount": "avg", + "*": "count", + } + ).rename({"l_returnflag": "al_returnflag", "l_linestatus": "al_linestatus"}) + + +# this doesn't seem like the right thing: +# def l_quantity_sum(grouped_lineitem: GroupedDataStream) -> pl.Series: +# pass + + +# hack to get around `@tag_outputs` not working with `@extract_columns` as expected. +# @extract_columns( +# *[ +# "__count_sum", +# "l_quantity_sum", +# "l_extendedprice_sum", +# "disc_price_sum", +# "charge_sum", +# "l_quantity_mean", +# "l_extendedprice_mean", +# "l_discount_mean", +# "al_returnflag", +# "al_linestatus", +# ] +# ) +# def extract_aggregates(compute_aggregates: DataStream) -> DataStream: +# return compute_aggregates + + +# tpc-h 3 +# on=None, left_on=None, right_on=None, suffix="_2", how="inner" +@tag(operation="join", left_on="l_orderkey", right_on="o_orderkey") +def lineitem_orders_table( + l_discount: pl.Series, + l_extendedprice: pl.Series, + l_orderkey: pl.Series, # what if there are two with same name? source__name syntax to disambiguate? + l_shipdate: pl.Series, + o_orderdate: pl.Series, + o_orderkey: pl.Series, + o_shippriority: pl.Series, +) -> DataStream: + pass + + +@extract_columns( + *[ + "lo_custkey", + "lo_orderkey", + "lo_orderdate", + "lo_shippriority", + "lo_shipdate", + "lo_extendedprice", + "lo_discount", + ] +) +def extract_lineitem_orders_table(lineitem_orders_table: DataStream) -> DataStream: + print(lineitem_orders_table.schema) + renamed_schema = { + col_name: f"lo_{'_'.join(col_name.split('_')[1:])}" + for col_name in lineitem_orders_table.schema + } + return lineitem_orders_table.rename(renamed_schema) + + +# TODO: handle joining on datastream object. +@tag( + operation="join", + left_on="c_custkey", + right_on="lo_custkey", + # TODO handle filtering...? + # fitler="c_mktsegment = 'BUILDING' and lo_orderdate < date '1995-03-15' and lo_shipdate > date '1995-03-15'" +) +def customer_lineitem_orders_table( + c_custkey: pl.Series, + c_mktsegment: pl.Series, + lo_custkey: pl.Series, + lo_orderdate: pl.Series, + lo_orderkey: pl.Series, + lo_shippriority: pl.Series, + lo_shipdate: pl.Series, +) -> DataStream: + pass + + +@extract_columns( + *[ + "clo_custkey", + "clo_orderkey", + "clo_orderdate", + "clo_shippriority", + "clo_shipdate", + "clo_extendedprice", + "clo_discount", + ] +) +def extract_customer_lineitem_orders_table( + customer_lineitem_orders_table: DataStream, +) -> DataStream: + print(customer_lineitem_orders_table.schema) + renamed_schema = { + col_name: f"clo_{'_'.join(col_name.split('_')[1:])}" + for col_name in customer_lineitem_orders_table.schema + } + return customer_lineitem_orders_table.rename(renamed_schema) + + +@tag(group_by="clo_orderkey,clo_orderdate,clo_shippriority") +def grouped_customer_lineitem_orders_table( + clo_orderkey: pl.Series, + clo_orderdate: pl.Series, + clo_shippriority: pl.Series, + clo_revenue: pl.Series, +) -> GroupedDataStream: + """This function declares the "schema" the datastream needs to have via it's arguments. + The body is blank because the graph adapter knows the actual logic to perform. + + Alternate syntax could be to have the decorator declare what's required, the function then takes the datastream, the + group by happens in the function... + """ + pass + + +@extract_columns(*["aclo_revenue_sum", "aclo_orderkey", "aclo_orderdate", "aclo_shippriority"]) +def aggregated_grouped_customer_lineitem_orders_table( + grouped_customer_lineitem_orders_table: GroupedDataStream, +) -> DataStream: + # Thought: change these into individual functions? + return grouped_customer_lineitem_orders_table.agg({"clo_revenue": "sum"}).rename( + { + "clo_revenue_sum": "aclo_revenue_sum", + "clo_orderkey": "aclo_orderkey", + "clo_orderdate": "aclo_orderdate", + "clo_shippriority": "aclo_shippriority", + } + ) diff --git a/examples/quokka/quokka_adapter.py b/examples/quokka/quokka_adapter.py new file mode 100644 index 00000000..6e1451b0 --- /dev/null +++ b/examples/quokka/quokka_adapter.py @@ -0,0 +1,175 @@ +import inspect +from typing import Any, Callable, Dict, Tuple + +import polars as pl +from pyquokka.datastream import DataStream, GroupedDataStream + +from hamilton import base, node + + +class QuokkaGraphAdapter(base.SimplePythonDataFrameGraphAdapter): + def __init__(self, result_builder: base.ResultMixin = base.DictResult()): + self.ds_objects = {} + self.call_count = 0 + self.result_builder = result_builder # not used... + + def _lambda_udf(self, ds: DataStream, udf: Callable) -> DataStream: + """Function to wrap pulling metadata from a hamilton function""" + sig = inspect.signature(udf) + input_parameters = dict(sig.parameters) + + def hack_wrapper(x): + # this is required because that's the signature datastream.with_column expects. i.e. a single arg function. + kwargs = {k: x[k] for k in input_parameters} + return udf(**kwargs) + + return ds.with_column( + udf.__name__, hack_wrapper, required_columns=set(input_parameters.keys()) + ) + + def _sanitize_kwargs(self, kwargs: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Sanitizes the kwargs to remove the datastream node name.""" + ds_names = {} + actual_kwargs = {} + for kwarg_key, kwarg_value in kwargs.items(): + if isinstance(kwarg_value, dict) and "__ds_name__" in kwarg_value: + ds_names[kwarg_key] = kwarg_value["__ds_name__"] + actual_kwargs[kwarg_key] = kwarg_value["__result__"] + else: + actual_kwargs[kwarg_key] = kwarg_value + return actual_kwargs, ds_names + + def execute_node(self, node: node.Node, kwargs: Dict[str, Any]) -> Any: + self.call_count += 1 + actual_kwargs, ds_names = self._sanitize_kwargs(kwargs) + ds_name_set = set(ds_names.values()) + if node.type == DataStream: + if node.tags.get("operation", None) == "join": + left = node.tags["left_on"] + right = node.tags["right_on"] + left_ds = self.ds_objects[ds_names[left]] + right_ds = self.ds_objects[ds_names[right]] + ds: DataStream = left_ds.join( + right_ds, left_on=left, right_on=right + ) # only do inner for now. + if node.tags.get("filter", None): + print(node.tags["filter"]) + ds = ds.filter(node.tags["filter"]) + self.ds_objects[node.name] = ds + return {"__ds_name__": node.name, "__result__": ds} + else: + # assumption is types into this function are only scalars, or other DataStream/GroupedDataStream objects + ds: DataStream = node.callable(**actual_kwargs) + self.ds_objects[node.name] = ds + return {"__ds_name__": node.name, "__result__": ds} + elif node.type == GroupedDataStream: + print("got group by", self.call_count, node.name) + assert len(ds_name_set) == 1, f"Error groupby got multiple DataStreams {ds_names}" + ds = self.ds_objects[ds_names.popitem()[1]] + print("before select", ds.schema) + ds = ds.select(list(node.input_types.keys())) + print("after select", ds.schema) + group_by_cols = node.tags["group_by"].split(",") + order_by_cols = node.tags["order_by"].split(",") if "order_by" in node.tags else None + ds: GroupedDataStream = ds.groupby(group_by_cols, orderby=order_by_cols) + self.ds_objects[node.name] = ds + return {"__ds_name__": node.name, "__result__": ds} + elif ( + node.type == pl.Series + and len(node.input_types) == 1 + and node.tags.get("__generated_by__", None) == "extract_columns" + ): + assert ( + len(ds_name_set) == 1 + ), f"Error extract_columns got multiple DataStreams {ds_names}" + ds_name = ds_name_set.pop() + print(node.name, node.tags) + # print('got extract_columns', self.call_count, node.name, kwargs) + # get global one + ds = self.ds_objects[ds_name] + return {"__ds_name__": ds_name, "__result__": ds} + else: + assert len(ds_name_set) == 1, f"Error udf got multiple DataStreams {ds_names}" + print("got udf", self.call_count, node.name, kwargs) + ds_name = ds_name_set.pop() + ds: DataStream = self.ds_objects[ds_name] + print(ds.schema) + ds = self._lambda_udf(ds, node.callable) + self.ds_objects[ds_name] = ds + return {"__ds_name__": ds_name, "__result__": ds} + + def execute_node2(self, node: node.Node, kwargs: Dict[str, Any]) -> Any: + """Very crappy code -- this is just to get something working. This should be refactored to be more generic.""" + """ + Idea to handle multiple input data streams + - DFS happens from outputs. So we should never have multiple DS's live as we progress through the DAG + (assuming no parallel branches that never meet up together) + - so what we need to keep track of is the name of the function that produced the datastream + (do we pass it through kwargs as a hack?) + - and who are the descendents of that function, so that they get the right datastream object. + - things reset/are replaced when either a groupby, aggregation, or join happen. + """ + if node.type == DataStream: + print("got datastream", self.call_count, node.name) + ds: DataStream = node.callable(**kwargs) + if node.tags.get("materialize", "False") == "True": + print("materializing", node.name) + ds = ds.collect() + self.call_count += 1 + return ds + elif node.type == GroupedDataStream: + print("got group by", self.call_count, node.name) + ds = self.ds_objects["global"] + print("before select", ds.schema) + ds = ds.select(list(node.input_types.keys())) + print("after select", ds.schema) + group_by_cols = node.tags["group_by"].split(",") + order_by_cols = node.tags["order_by"].split(",") + ds: GroupedDataStream = ds.groupby(group_by_cols, orderby=order_by_cols) + self.call_count += 1 + return ds + elif node.type == pl.Series and len(node.input_types) == 1: + # print('got extract_columns', self.call_count, node.name, kwargs) + # get global one + ds = self.ds_objects["global"] + else: + print("got udf", self.call_count, node.name, kwargs) + ds: DataStream = self.ds_objects["global"] + print(ds.schema) + ds = self._lambda_udf(ds, node.callable) + self.ds_objects["global"] = ds + self.call_count += 1 + return ds + + def build_result(self, **outputs: Dict[str, Any]) -> pl.DataFrame: + """Builds the result and brings it back to this running process. + + :param outputs: the dictionary of key -> Union[ray object reference | value] + :return: The type of object returned by self.result_builder. + """ + # TODO: this is a bit hacky. `Collect()` should ideally be only called here. + # Right now we assume that the final global object is what we need. + # also this is brittle and would break if we request intermediate nodes for example. + requested_ds_set = set(outputs.keys()) + actual_outputs, ds_names = self._sanitize_kwargs(outputs) + ds_name_set = set(ds_names.values()) + assert ( + len(ds_name_set) == 1 + ), f"Error got multiple DataStreams to build result from {ds_names}" + ds = self.ds_objects[ds_name_set.pop()] + if isinstance(ds, (DataStream, GroupedDataStream)): + print(ds.explain(mode="text")) + df = ds.collect() + # if isinstance(ds, pl.DataFrame): + df = df[list(actual_outputs.keys())] + # else: + # hack to get around issue with schema not being passed through correctly + # df = ds.collect() + # df: pl.DataFrame = ds.select(list(actual_outputs.keys())).collect() + global_ds_set = set(df.columns) + if requested_ds_set.intersection(global_ds_set) != requested_ds_set: + raise ValueError( + f"Error: requested columns not found in final dataframe. " + f"Missing: {requested_ds_set.difference(global_ds_set)}." + ) + return df diff --git a/examples/quokka/run_adapter.py b/examples/quokka/run_adapter.py new file mode 100644 index 00000000..ac38812b --- /dev/null +++ b/examples/quokka/run_adapter.py @@ -0,0 +1,55 @@ +import possible_api +import quokka_adapter +from pyquokka.df import QuokkaContext + +from hamilton import driver + + +def tpc_1(qc, path): + adapter = quokka_adapter.QuokkaGraphAdapter() + dr = driver.Driver( + { + "qc": qc, + "path": path, + "lineitem_filter": "l_shipdate <= date '1998-12-01' - interval '90' day", + }, + possible_api, + adapter=adapter, + ) + outputs = [ + "al_returnflag", # can comment one of these out and they are dropped + "al_linestatus", + "__count_sum", + "charge_sum", + "disc_price_sum", + "l_discount_mean", + "l_extendedprice_mean", + "l_extendedprice_sum", + "l_quantity_mean", + "l_quantity_sum", + # "grouped_lineitem" -- calling collect on groups doesn't work it seems... + # "l_returnflag", # can also get intermediate results + # "charge", + # "disc_price", + ] + dr.visualize_execution(outputs, "./tpc_1.dot", {}) + result = dr.execute(outputs) + print(result.columns) + print(result) + + +def tpc_3(qc, path): + adapter = quokka_adapter.QuokkaGraphAdapter() + dr = driver.Driver({"qc": qc, "path": path}, possible_api, adapter=adapter) + outputs = ["aclo_orderkey", "aclo_orderdate", "aclo_shippriority", "aclo_revenue_sum"] + dr.visualize_execution(outputs, "./tpc_3.dot", {}) + result = dr.execute(outputs) + print(result.columns) + print(result) + + +if __name__ == "__main__": + qc = QuokkaContext() + path = "/Users/stefankrawczyk/Downloads/tpc-h-public/" + # tpc_1(qc, path) + tpc_3(qc, path) diff --git a/hamilton/function_modifiers/expanders.py b/hamilton/function_modifiers/expanders.py index c1372c51..b9824c5f 100644 --- a/hamilton/function_modifiers/expanders.py +++ b/hamilton/function_modifiers/expanders.py @@ -307,6 +307,23 @@ class parameterized_inputs(parameterize_sources): pass +# HACKS to get quokka working --- DO NOT MERGE! +try: + import polars as pl + from pyquokka import datastream as qds + + branch = "quokka" +except ImportError: + pass +try: + from pyspark.sql import Column as SparkColumn + from pyspark.sql import DataFrame as SparkDataFrame + + branch = "pyspark" +except ImportError: + pass + + class extract_columns(base.NodeExpander): def __init__(self, *columns: Union[Tuple[str, str], str], fill_with: Any = None): """Constructor for a modifier that expands a single function into the following nodes: @@ -335,10 +352,11 @@ def validate(self, fn: Callable): :raises: InvalidDecoratorException If the function does not output a Dataframe """ output_type = inspect.signature(fn).return_annotation - if not issubclass(output_type, pd.DataFrame): - raise base.InvalidDecoratorException( - f"For extracting columns, output type must be pandas dataframe, not: {output_type}" - ) + print(output_type) + # if not issubclass(output_type, pd.DataFrame): + # raise base.InvalidDecoratorException( + # f"For extracting columns, output type must be pandas dataframe, not: {output_type}" + # ) def expand_node( self, node_: node.Node, config: Dict[str, Any], fn: Callable @@ -356,10 +374,10 @@ def expand_node( def df_generator(*args, **kwargs) -> pd.DataFrame: df_generated = fn(*args, **kwargs) - if self.fill_with is not None: - for col in self.columns: - if col not in df_generated: - df_generated[col] = self.fill_with + # if self.fill_with is not None: + # for col in self.columns: + # if col not in df_generated: + # df_generated[col] = self.fill_with return df_generated output_nodes = [node_.copy_with(callabl=df_generator)] @@ -373,21 +391,40 @@ def extractor_fn( column_to_extract: str = column, **kwargs ) -> pd.Series: # avoiding problems with closures df = kwargs[node_.name] - if column_to_extract not in df: + # print(node_.name, column_to_extract, type(df)) + if column_to_extract not in df.columns: raise base.InvalidDecoratorException( f"No such column: {column_to_extract} produced by {node_.name}. " f"It only produced {str(df.columns)}" ) - return kwargs[node_.name][column_to_extract] + if branch == "pyspark" and isinstance(df, SparkDataFrame): + return df # df.select(column_to_extract) # this returns a SparkDataFrame not a column... but yeah + + if branch == "quokka" and isinstance(df, pl.DataFrame): + return df[column_to_extract] + return df + + if branch == "quokka": + output_type = qds.DataStream + input_type = pl.Series + elif branch == "pyspark": + output_type = SparkColumn + input_type = SparkDataFrame + else: + raise ValueError(f"unknow branch {branch}") + tags = node_.tags.copy() + tags["__generated_by__"] = "extract_columns" output_nodes.append( node.Node( column, - pd.Series, + input_type, # SparkColumn, # pl.Series, # pd.Series, doc_string, extractor_fn, - input_types={node_.name: pd.DataFrame}, - tags=node_.tags.copy(), + input_types={ + node_.name: output_type + }, # SparkDataFrame}, # qds.DataStream}, # pd.DataFrame}, + tags=tags, ) ) return output_nodes