@@ -17,7 +17,6 @@ const MIN_RETRY_DELAY: Duration = Duration::from_secs(1);
1717const MAX_RETRY_DELAY : Duration = Duration :: from_secs ( 16 ) ;
1818
1919/// Everything the downloader needs to fetch and parse a single remote partition.
20- #[ derive( Clone ) ]
2120pub ( super ) struct DownloadRequest {
2221 pub ( super ) url : String ,
2322 pub ( super ) headers : Arc < HeaderMap > ,
@@ -34,6 +33,11 @@ enum DownloadFailure {
3433 Fatal ( QueryScopedError ) ,
3534}
3635
36+ struct DownloadedChunk {
37+ body : Bytes ,
38+ workload : ParseWorkload ,
39+ }
40+
3741impl DownloadFailure {
3842 fn into_inner ( self ) -> QueryScopedError {
3943 match self {
@@ -59,12 +63,19 @@ impl RemotePartitionDownloader {
5963 schema : Arc < Schema > ,
6064 ) -> QueryScopedResult < ResultTable > {
6165 let mut retries = 0 ;
66+ let column_count = schema. len ( ) ;
6267 loop {
63- match self
64- . download_once ( request. clone ( ) , Arc :: clone ( & schema) )
65- . await
66- {
67- Ok ( t) => return Ok ( t) ,
68+ match self . download_once ( & request, column_count) . await {
69+ Ok ( chunk) => {
70+ return parse_remote_chunk_result_table_async (
71+ schema,
72+ request. query_id ,
73+ chunk. body ,
74+ chunk. workload ,
75+ request. blocking_parse_limiter ,
76+ )
77+ . await ;
78+ }
6879 Err ( DownloadFailure :: Retryable ( _) ) if retries < MAX_RETRIES => {
6980 sleep ( retry_delay ( retries) ) . await ;
7081 retries += 1 ;
@@ -76,9 +87,9 @@ impl RemotePartitionDownloader {
7687
7788 async fn download_once (
7889 & self ,
79- request : DownloadRequest ,
80- schema : Arc < Schema > ,
81- ) -> StdResult < ResultTable , DownloadFailure > {
90+ request : & DownloadRequest ,
91+ column_count : usize ,
92+ ) -> StdResult < DownloadedChunk , DownloadFailure > {
8293 let response = match self
8394 . client
8495 . get ( & request. url )
@@ -90,12 +101,12 @@ impl RemotePartitionDownloader {
90101 Err ( error) => {
91102 if error. is_timeout ( ) {
92103 return Err ( DownloadFailure :: Retryable ( QueryScopedError :: new (
93- request. query_id ,
104+ Arc :: clone ( & request. query_id ) ,
94105 TimeoutError :: request ( error) ,
95106 ) ) ) ;
96107 }
97108 return Err ( DownloadFailure :: Retryable ( QueryScopedError :: new (
98- request. query_id ,
109+ Arc :: clone ( & request. query_id ) ,
99110 NetworkError :: Http ( error) ,
100111 ) ) ) ;
101112 }
@@ -106,7 +117,7 @@ impl RemotePartitionDownloader {
106117 return Err ( classify_failed_status (
107118 status,
108119 response. bytes ( ) . await ,
109- request. query_id ,
120+ Arc :: clone ( & request. query_id ) ,
110121 ) ) ;
111122 }
112123
@@ -115,12 +126,12 @@ impl RemotePartitionDownloader {
115126 Err ( error) => {
116127 if error. is_timeout ( ) {
117128 return Err ( DownloadFailure :: Retryable ( QueryScopedError :: new (
118- request. query_id ,
129+ Arc :: clone ( & request. query_id ) ,
119130 TimeoutError :: request ( error) ,
120131 ) ) ) ;
121132 }
122133 return Err ( DownloadFailure :: Retryable ( QueryScopedError :: new (
123- request. query_id ,
134+ Arc :: clone ( & request. query_id ) ,
124135 NetworkError :: Http ( error) ,
125136 ) ) ) ;
126137 }
@@ -130,21 +141,13 @@ impl RemotePartitionDownloader {
130141 let workload = ParseWorkload :: remote_chunk (
131142 body. len ( ) ,
132143 request. row_count ,
133- schema . len ( ) ,
144+ column_count ,
134145 gzip_encoded,
135146 request. compressed_size ,
136147 request. uncompressed_size ,
137148 ) ;
138149
139- parse_remote_chunk_result_table_async (
140- schema,
141- Arc :: clone ( & request. query_id ) ,
142- body,
143- workload,
144- request. blocking_parse_limiter ,
145- )
146- . await
147- . map_err ( DownloadFailure :: Fatal )
150+ Ok ( DownloadedChunk { body, workload } )
148151 }
149152}
150153
0 commit comments