1+ import pandas as pd
2+ import polars as pl
3+ import sqlalchemy as sa
4+ import ibis
5+ from ibis import _ as col
6+ import pydiverse .transform as pdt
7+ from pydiverse .pipedag import Flow , Stage , Table , materialize
8+ from pydiverse .transform import λ
9+ from pydiverse .transform .core .functions import row_number
10+ from pydiverse .transform .core .verbs import (
11+ mutate , alias , arrange , build_query ,
12+ )
13+ from pydiverse .transform .eager import PandasTableImpl
14+ from pydiverse .transform .lazy import SQLTableImpl
15+
16+
17+ @materialize (version = "1.0.0" )
18+ def read_input_data ():
19+ titanic = pd .read_csv (
20+ 'https://raw.githubusercontent.com/mwaskom/seaborn-data/master/titanic.csv'
21+ )
22+ return Table (titanic , name = "titanic" )
23+
24+
25+ @materialize (input_type = pd .DataFrame , version = "win1.0.0" )
26+ def task_pandas (titanic : pd .DataFrame ):
27+ titanic .sort_values ("fare" , inplace = True )
28+ return (
29+ titanic
30+ .assign (idx = range (len (titanic )), diff_price = titanic .fare .diff ())
31+ )
32+
33+
34+ @materialize (input_type = pl .DataFrame , version = "win1.0.0" )
35+ def task_polars (titanic : pl .DataFrame ):
36+ return (
37+ titanic .sort ("fare" )
38+ .with_columns (idx = pl .Series (range (len (titanic ))), diff_price = pl .col ("fare" ).diff ())
39+ )
40+
41+
42+ @materialize (input_type = PandasTableImpl , version = "win1.0.1" )
43+ def task_transform_df (titanic : pdt .Table ):
44+ return (
45+ titanic
46+ >> mutate (
47+ idx = row_number (arrange = [λ .fare ]),
48+ diff_price = λ .fare .shift (- 1 , arrange = [λ .fare ])- λ .fare
49+ )
50+ >> arrange (λ .fare )
51+ >> alias ("transform_df" )
52+ )
53+
54+
55+ @materialize (input_type = SQLTableImpl , lazy = True )
56+ def task_transform_sql (titanic : pdt .Table ):
57+ return (
58+ titanic
59+ >> mutate (
60+ idx = row_number (arrange = [λ .fare ]),
61+ diff_price = λ .fare .shift (- 1 , arrange = [λ .fare ])- λ .fare
62+ )
63+ >> arrange (λ .fare )
64+ >> alias ("transform_sql" )
65+ )
66+
67+
68+ @materialize (input_type = ibis .api .Table , lazy = True )
69+ def task_ibis (titanic : ibis .api .Table ):
70+ w = ibis .window (order_by = col .fare )
71+ return (
72+ titanic
73+ .mutate (idx = ibis .row_number ().over (w ), diff_price = col .fare .lag (- 1 ).over (w ) - col .fare )
74+ .order_by (col .fare )
75+ )
76+
77+
78+ @materialize (input_type = sa .Table , lazy = True )
79+ def task_sqlalchemy (titanic : sa .Table ):
80+ return sa .select (
81+ titanic ,
82+ sa .func .row_number ().over (order_by = titanic .c .fare ).label ("idx" ),
83+ (sa .func .lead (titanic .c .fare , 1 , None ).over (order_by = titanic .c .fare )
84+ - titanic .c .fare ).label ("diff_price" ),
85+ ).select_from (titanic ).order_by (titanic .c .fare )
86+
87+
88+ @materialize (input_type = sa .Table , lazy = True )
89+ def task_sql (titanic : sa .Table ):
90+ return sa .text (f"""
91+ SELECT tt.*,
92+ ROW_NUMBER() OVER (ORDER BY tt.fare ASC NULLS LAST) AS idx,
93+ LEAD(tt.fare, 1, NULL) OVER (ORDER BY tt.fare ASC NULLS LAST)
94+ - tt.fare AS diff_price
95+ FROM { titanic .original .schema } { titanic .name }
96+ ORDER BY tt.fare
97+ """ )
98+
99+
100+ @materialize (input_type = pd .DataFrame )
101+ def print_tables (tbls : list [pd .DataFrame ]):
102+ from matplotlib import pyplot as plt
103+ fix , axs = plt .subplots (nrows = 4 , ncols = 2 )
104+ for tbl , ax in zip (tbls , axs .flatten ()):
105+ tbl .sort_values ("idx" , inplace = True )
106+ print (f"\n \n { tbl } )
107+ ax .plot (tbl .idx , tbl .diff_price .fillna (0 ))
108+ # limit y axis to 30
109+ ax .set_ylim (0 , 30 )
110+ plt .show ()
111+
112+
113+ def get_pipeline ():
114+ tasks = [task_pandas , task_polars , task_transform_df , task_transform_sql ,
115+ task_ibis , task_sqlalchemy , task_sql ]
116+ with Flow ("flow" ) as flow :
117+ with Stage ("t1_raw_input" ):
118+ titanic = read_input_data ()
119+
120+ with Stage ("t2_transformed_data" ):
121+ out_tbls = [task (titanic ) for task in tasks ]
122+ print_tables (out_tbls )
123+
124+ return flow
125+
126+
127+ if __name__ == "__main__" :
128+ import logging
129+ from pydiverse .pipedag .util .structlog import setup_logging
130+
131+ setup_logging (log_level = logging .INFO )
132+
133+ flow = get_pipeline ()
134+ result = flow .run ()
135+ assert result .successful
136+ flow .run (flow ["t2_transformed_data" ]["task_transform_sql" ])
0 commit comments