@@ -587,11 +587,11 @@ mod tests {
587587 #[ test]
588588 fn selected_candidate_input_budget_is_a_terminal_error ( ) {
589589 let plan = CompiledPhysicalDag :: decode ( & sorted ( ) ) . unwrap ( ) ;
590- let error = execute_batches ( & plan, 1 , 1 , 42 , |_, schema | {
591- Ok ( Batch :: try_new (
592- schema. clone ( ) ,
590+ let error = execute_batches ( & plan, 1 , 1 , 42 , |_, contract | {
591+ Ok ( BoundInput :: Rows ( Batch :: try_new (
592+ contract . schema . clone ( ) ,
593593 vec ! [ vec![ Value :: Float64 ( 1. ) ] ] ,
594- ) ?)
594+ ) ?) )
595595 } )
596596 . expect_err ( "input must exceed the byte budget" ) ;
597597 assert ! (
@@ -767,7 +767,8 @@ where
767767 max_bytes,
768768 bindings. len ( ) ,
769769 at,
770- |input_id, schema| {
770+ |input_id, contract| {
771+ let schema = & contract. schema ;
771772 let values = super :: vector ( super :: from_result ( callback ( bindings[ & input_id] , at) ?) ?) ?;
772773 let rows = values
773774 . into_iter ( )
@@ -781,7 +782,9 @@ where
781782 . map_err ( |e| miss ( e. to_string ( ) ) )
782783 } )
783784 . collect :: < Result < Vec < _ > , _ > > ( ) ?;
784- Batch :: try_new ( schema. clone ( ) , rows) . map_err ( EngineError :: from)
785+ Batch :: try_new ( schema. clone ( ) , rows)
786+ . map ( BoundInput :: Rows )
787+ . map_err ( EngineError :: from)
785788 } ,
786789 )
787790}
@@ -790,7 +793,8 @@ pub(in crate::query_engines::asap_query_engine) fn execute_stored(
790793 entry : & asap_types:: query_plan:: QueryPlanEntry ,
791794 plan_id : u64 ,
792795 plan_version : u64 ,
793- store : & crate :: storage_engines:: sketch_db:: index:: SketchStore ,
796+ store : Option < & crate :: storage_engines:: sketch_db:: index:: SketchStore > ,
797+ raw_endpoint : Option < ( & reqwest:: Client , & str ) > ,
794798 at : u64 ,
795799) -> Result <
796800 (
@@ -805,16 +809,28 @@ pub(in crate::query_engines::asap_query_engine) fn execute_stored(
805809 let program = entry
806810 . recover_vector_physical_dag ( )
807811 . map_err ( |e| miss ( e. to_string ( ) ) ) ?;
808- execute_batches ( & program, max_bytes, inputs. len ( ) , at, |id, schema| {
812+ execute_batches ( & program, max_bytes, inputs. len ( ) , at, |id, contract| {
813+ let schema = & contract. schema ;
809814 let index = sources
810815 . iter ( )
811816 . position ( |source| * source == id)
812817 . ok_or_else ( || miss ( "native source is unbound" ) ) ?;
813- let Some ( asap_types:: query_plan:: QueryPlanNode :: ReadMaterialization { binding } ) =
814- entry. nodes . get ( & inputs[ index] )
818+ let node = entry. nodes . get ( & inputs[ index] ) ;
819+ if let Some ( asap_types:: query_plan:: QueryPlanNode :: Logical { operator, .. } ) = node {
820+ let ( client, endpoint) = raw_endpoint
821+ . ok_or_else ( || miss ( "query-time raw input has no Prometheus endpoint" ) ) ?;
822+ let at = i64:: try_from ( at) . map_err ( |_| miss ( "native timestamp overflow" ) ) ?;
823+ return Ok ( BoundInput :: Lazy ( super :: super :: raw_source:: bind (
824+ contract, operator, at, client, endpoint,
825+ ) ?) ) ;
826+ }
827+ let Some ( asap_types:: query_plan:: QueryPlanNode :: ReadMaterialization { binding } ) = node
815828 else {
816829 return Err ( miss ( "native stored source has no deployed summary binding" ) ) ;
817830 } ;
831+ let store = store. ok_or_else ( || {
832+ EngineError :: capability_miss ( "native_stored" , "summary store unavailable" )
833+ } ) ?;
818834 let end = i64:: try_from ( at) . map_err ( |_| miss ( "native timestamp overflow" ) ) ?;
819835 let start = at
820836 . checked_sub ( binding. window_ms )
@@ -836,6 +852,7 @@ pub(in crate::query_engines::asap_query_engine) fn execute_stored(
836852 schema. clone ( ) ,
837853 max_bytes as usize ,
838854 )
855+ . map ( BoundInput :: Rows )
839856 . map_err ( |error| match error {
840857 crate :: storage_engines:: sketch_db:: index:: NativeReadError :: Unavailable ( message) => {
841858 miss ( message)
@@ -847,12 +864,22 @@ pub(in crate::query_engines::asap_query_engine) fn execute_stored(
847864 } )
848865}
849866
867+ /// Stored and protocol inputs are read before execution; query-time raw
868+ /// sources open lazily inside the run and charge its budget as they stream.
869+ enum BoundInput {
870+ Rows ( Batch ) ,
871+ Lazy ( asap_physical_operators:: physical_planner:: Source < ' static > ) ,
872+ }
873+
850874fn execute_batches (
851875 program : & asap_physical_operators:: physical_planner:: CompiledPhysicalDag ,
852876 max_bytes : u64 ,
853877 input_count : usize ,
854878 at : u64 ,
855- mut input_batch : impl FnMut ( u64 , & Schema ) -> Result < Batch , EngineError > ,
879+ mut input_batch : impl FnMut (
880+ u64 ,
881+ & asap_physical_operators:: physical_planner:: InputContract ,
882+ ) -> Result < BoundInput , EngineError > ,
856883) -> Result <
857884 (
858885 crate :: query_engines:: query_result:: QueryResult ,
@@ -887,7 +914,13 @@ fn execute_batches(
887914 let mut input_bytes = 0usize ;
888915 for ( input_id, input) in program. input_contracts ( ) {
889916 crate :: query_engines:: request:: check ( ) ?;
890- let batch = input_batch ( input_id, & input. schema ) ?;
917+ let batch = match input_batch ( input_id, input) ? {
918+ BoundInput :: Rows ( batch) => batch,
919+ BoundInput :: Lazy ( source) => {
920+ sources. insert ( input_id, source) ;
921+ continue ;
922+ }
923+ } ;
891924 input_bytes = input_bytes
892925 . checked_add ( batch. bytes ( ) )
893926 . ok_or ( asap_physical_operators:: Error :: MemoryLimit ) ?;
@@ -996,6 +1029,7 @@ mod request_contract_tests {
9961029 ) ?;
9971030 execute_batches ( & program, 64 * 1024 , 1 , 0 , |_, _| {
9981031 Batch :: try_new ( schema. clone ( ) , vec ! [ vec![ Value :: Float64 ( 1.0 ) ] ] )
1032+ . map ( BoundInput :: Rows )
9991033 . map_err ( EngineError :: from)
10001034 } )
10011035 } ,
0 commit comments