13
13
// limitations under the License.
14
14
15
15
use std:: any:: Any ;
16
- use std:: collections:: VecDeque ;
17
16
use std:: str:: FromStr ;
18
17
use std:: sync:: Arc ;
19
18
@@ -29,15 +28,14 @@ use common_legacy_planners::Statistics;
29
28
use common_meta_app:: schema:: TableInfo ;
30
29
use common_pipeline_core:: processors:: port:: InputPort ;
31
30
use common_pipeline_core:: SinkPipeBuilder ;
31
+ use common_pipeline_sources:: processors:: sources:: input_formats:: InputContext ;
32
32
use parking_lot:: Mutex ;
33
33
use tracing:: info;
34
34
35
35
use super :: StageSourceHelper ;
36
- use crate :: pipelines:: processors:: port:: OutputPort ;
37
36
use crate :: pipelines:: processors:: ContextSink ;
38
37
use crate :: pipelines:: processors:: TransformLimit ;
39
38
use crate :: pipelines:: Pipeline ;
40
- use crate :: pipelines:: SourcePipeBuilder ;
41
39
use crate :: sessions:: TableContext ;
42
40
use crate :: storages:: Table ;
43
41
@@ -47,6 +45,7 @@ pub struct StageTable {
47
45
// But the Table trait need it:
48
46
// fn get_table_info(&self) -> &TableInfo).
49
47
table_info_placeholder : TableInfo ,
48
+ input_context : Mutex < Option < Arc < InputContext > > > ,
50
49
}
51
50
52
51
impl StageTable {
@@ -56,8 +55,14 @@ impl StageTable {
56
55
Ok ( Arc :: new ( Self {
57
56
table_info,
58
57
table_info_placeholder,
58
+ input_context : Default :: default ( ) ,
59
59
} ) )
60
60
}
61
+
62
+ fn get_input_context ( & self ) -> Option < Arc < InputContext > > {
63
+ let guard = self . input_context . lock ( ) ;
64
+ guard. clone ( )
65
+ }
61
66
}
62
67
63
68
#[ async_trait:: async_trait]
@@ -73,39 +78,35 @@ impl Table for StageTable {
73
78
74
79
async fn read_partitions (
75
80
& self ,
76
- _ctx : Arc < dyn TableContext > ,
81
+ ctx : Arc < dyn TableContext > ,
77
82
_push_downs : Option < Extras > ,
78
83
) -> Result < ( Statistics , Partitions ) > {
84
+ let operator = StageSourceHelper :: get_op ( & ctx, & self . table_info . stage_info ) . await ?;
85
+ let input_ctx = Arc :: new (
86
+ InputContext :: try_create_from_copy (
87
+ operator,
88
+ ctx. get_settings ( ) . clone ( ) ,
89
+ ctx. get_format_settings ( ) ?,
90
+ self . table_info . schema . clone ( ) ,
91
+ self . table_info . stage_info . clone ( ) ,
92
+ self . table_info . files . clone ( ) ,
93
+ ctx. get_scan_progress ( )
94
+ )
95
+ . await ?,
96
+ ) ;
97
+ let mut guard = self . input_context . lock ( ) ;
98
+ * guard = Some ( input_ctx) ;
79
99
Ok ( ( Statistics :: default ( ) , vec ! [ ] ) )
80
100
}
81
101
82
102
fn read2 (
83
103
& self ,
84
- ctx : Arc < dyn TableContext > ,
104
+ _ctx : Arc < dyn TableContext > ,
85
105
_plan : & ReadDataSourcePlan ,
86
106
pipeline : & mut Pipeline ,
87
107
) -> Result < ( ) > {
88
- let settings = ctx. get_settings ( ) ;
89
- let mut builder = SourcePipeBuilder :: create ( ) ;
90
- let table_info = & self . table_info ;
91
- let schema = table_info. schema . clone ( ) ;
92
- let mut files_deque = VecDeque :: with_capacity ( table_info. files . len ( ) ) ;
93
- for f in & table_info. files {
94
- files_deque. push_back ( f. to_string ( ) ) ;
95
- }
96
- let files = Arc :: new ( Mutex :: new ( files_deque) ) ;
97
-
98
- let stage_source = StageSourceHelper :: try_create ( ctx, schema, table_info. clone ( ) , files) ?;
99
-
100
- for _index in 0 ..settings. get_max_threads ( ) ? {
101
- let output = OutputPort :: create ( ) ;
102
- builder. add_source ( output. clone ( ) , stage_source. get_splitter ( output) ?) ;
103
- }
104
- pipeline. add_pipe ( builder. finalize ( ) ) ;
105
-
106
- pipeline. add_transform ( |transform_input_port, transform_output_port| {
107
- stage_source. get_deserializer ( transform_input_port, transform_output_port)
108
- } ) ?;
108
+ let input_ctx = self . get_input_context ( ) . unwrap ( ) ;
109
+ input_ctx. format . exec_copy ( input_ctx. clone ( ) , pipeline) ?;
109
110
110
111
let limit = self . table_info . stage_info . copy_options . size_limit ;
111
112
if limit > 0 {
0 commit comments