@@ -98,6 +98,11 @@ def _build_task_registry(self) -> dict[str, callable]:
9898 # Join queries
9999 "inner_join" : self ._task_inner_join ,
100100 "left_join" : self ._task_left_join ,
101+ # Sort queries
102+ "sort_single" : self ._task_sort_single ,
103+ "sort_multi" : self ._task_sort_multi ,
104+ # Window join queries
105+ "window_join" : self ._task_window_join ,
101106 # Generic execution
102107 "sql" : self ._task_sql ,
103108 }
@@ -468,15 +473,87 @@ def _task_left_join(self, params: dict[str, Any]) -> AdapterResult:
468473 """Left join on id1, id2"""
469474 left_table = params .get ("left_table" , "x" )
470475 right_table = params .get ("right_table" , "y" )
471-
476+
472477 left_lf = self ._get_lazy_table (left_table )
473478 right_lf = self ._get_lazy_table (right_table )
474-
479+
475480 def query ():
476481 return left_lf .join (right_lf , on = ["id1" , "id2" ], how = "left" )
477-
482+
478483 return self ._execute_lazy (query , f'{ left_table } .join({ right_table } , on=["id1","id2"], how="left")' )
479-
484+
485+ # =========================================================================
486+ # Sort Queries (using lazy evaluation for optimal parallelism)
487+ # =========================================================================
488+
489+ def _task_sort_single (self , params : dict [str , Any ]) -> AdapterResult :
490+ """Sort by single column"""
491+ table_name = params .get ("table" , self ._table_name )
492+ column = params .get ("column" , "id1" )
493+ descending = params .get ("descending" , False )
494+ lf = self ._get_lazy_table (table_name )
495+
496+ def query ():
497+ return lf .sort (column , descending = descending )
498+
499+ return self ._execute_lazy (query , f'lf.sort("{ column } ", descending={ descending } )' )
500+
501+ def _task_sort_multi (self , params : dict [str , Any ]) -> AdapterResult :
502+ """Sort by multiple columns"""
503+ table_name = params .get ("table" , self ._table_name )
504+ columns = params .get ("columns" , ["id1" , "id2" ])
505+ lf = self ._get_lazy_table (table_name )
506+
507+ def query ():
508+ return lf .sort (columns )
509+
510+ return self ._execute_lazy (query , f'lf.sort({ columns } )' )
511+
512+ # =========================================================================
513+ # Window Join Queries (using join_asof for time-series joins)
514+ # =========================================================================
515+
516+ def _task_window_join (self , params : dict [str , Any ]) -> AdapterResult :
517+ """Window join - join within time window with aggregations
518+
519+ Polars doesn't have a direct wj1 equivalent, so we use a range join
520+ approach with group_by for aggregation.
521+ """
522+ import datetime
523+
524+ trades_table = params .get ("trades_table" , "trades" )
525+ quotes_table = params .get ("quotes_table" , "quotes" )
526+ window_ms = params .get ("window_ms" , 10000 ) # +/- 10 seconds default
527+
528+ trades_lf = self ._get_lazy_table (trades_table )
529+ quotes_lf = self ._get_lazy_table (quotes_table )
530+
531+ def query ():
532+ # Add window boundaries to trades
533+ trades_with_window = trades_lf .with_columns ([
534+ (pl .col ("Ts" ) - datetime .timedelta (milliseconds = window_ms )).alias ("_window_start" ),
535+ (pl .col ("Ts" ) + datetime .timedelta (milliseconds = window_ms )).alias ("_window_end" ),
536+ ])
537+
538+ # Cross join on Sym, filter by time window, then aggregate
539+ # This is expensive but semantically correct
540+ result = (
541+ trades_with_window
542+ .join (quotes_lf , on = "Sym" , how = "left" )
543+ .filter (
544+ (pl .col ("Ts_right" ) >= pl .col ("_window_start" )) &
545+ (pl .col ("Ts_right" ) <= pl .col ("_window_end" ))
546+ )
547+ .group_by (["Sym" , "Ts" , "Price" ])
548+ .agg ([
549+ pl .min ("Bid" ).alias ("Bid" ),
550+ pl .max ("Ask" ).alias ("Ask" ),
551+ ])
552+ )
553+ return result
554+
555+ return self ._execute_lazy (query , 'window_join(trades, quotes)' )
556+
480557 def _task_sql (self , params : dict [str , Any ]) -> AdapterResult :
481558 """Execute arbitrary SQL query."""
482559 query = params .get ("query" )
0 commit comments