1313// limitations under the License.
1414
1515use std:: collections:: BTreeMap ;
16+ use std:: sync:: atomic:: AtomicU64 ;
17+ use std:: sync:: atomic:: Ordering ;
1618use std:: sync:: Arc ;
1719
20+ use databend_common_base:: base:: tokio:: sync:: RwLock ;
1821use databend_common_exception:: Result ;
19- use databend_common_expression:: types:: DataType ;
2022use databend_common_expression:: types:: UInt64Type ;
2123use databend_common_expression:: DataBlock ;
2224use databend_common_expression:: FromData ;
@@ -30,34 +32,181 @@ use crate::sessions::QueryContext;
3032use crate :: sql:: executor:: physical_plans:: AsyncFunctionDesc ;
3133use crate :: sql:: plans:: AsyncFunctionArgument ;
3234
35+ // Structure to manage sequence numbers in batches
36+ pub struct SequenceCounter {
37+ // Current sequence number
38+ current : AtomicU64 ,
39+ // Maximum sequence number in the current batch
40+ max : AtomicU64 ,
41+ }
42+
43+ impl SequenceCounter {
44+ fn new ( ) -> Self {
45+ Self {
46+ current : AtomicU64 :: new ( 0 ) ,
47+ max : AtomicU64 :: new ( 0 ) ,
48+ }
49+ }
50+
51+ // Try to reserve a range of sequence numbers
52+ fn try_reserve ( & self , count : u64 ) -> Option < ( u64 , u64 ) > {
53+ if self . current . load ( Ordering :: Relaxed ) == 0 {
54+ return None ;
55+ }
56+
57+ let current = self . current . load ( Ordering :: Relaxed ) ;
58+ let max = self . max . load ( Ordering :: Relaxed ) ;
59+
60+ // Check if we have enough sequence numbers in the current batch
61+ if current + count <= max {
62+ let new_current = current + count;
63+ if self
64+ . current
65+ . compare_exchange ( current, new_current, Ordering :: SeqCst , Ordering :: Relaxed )
66+ . is_ok ( )
67+ {
68+ // Successfully reserved the range
69+ return Some ( ( current, new_current) ) ;
70+ }
71+ }
72+
73+ // Failed to reserve
74+ None
75+ }
76+
77+ // Update the counter with a new batch of sequence numbers
78+ fn update_batch ( & self , start : u64 , count : u64 ) {
79+ self . current . store ( start, Ordering :: SeqCst ) ;
80+ self . max . store ( start + count, Ordering :: SeqCst ) ;
81+ }
82+ }
83+
84+ // Shared sequence counters type
85+ pub type SequenceCounters = Vec < Arc < RwLock < SequenceCounter > > > ;
86+
3387pub struct TransformAsyncFunction {
3488 ctx : Arc < QueryContext > ,
3589 // key is the index of async_func_desc
3690 pub ( crate ) operators : BTreeMap < usize , Arc < DictionaryOperator > > ,
3791 async_func_descs : Vec < AsyncFunctionDesc > ,
92+ // Shared map of sequence name to sequence counter
93+ pub ( crate ) sequence_counters : SequenceCounters ,
3894}
3995
4096impl TransformAsyncFunction {
97+ // New constructor that accepts a shared sequence counters map
4198 pub ( crate ) fn new (
4299 ctx : Arc < QueryContext > ,
43100 async_func_descs : Vec < AsyncFunctionDesc > ,
44101 operators : BTreeMap < usize , Arc < DictionaryOperator > > ,
102+ sequence_counters : SequenceCounters ,
45103 ) -> Self {
46104 Self {
47105 ctx,
48106 async_func_descs,
49107 operators,
108+ sequence_counters,
50109 }
51110 }
52111
112+ // Create a new shared sequence counters map
113+ pub ( crate ) fn create_sequence_counters ( size : usize ) -> SequenceCounters {
114+ ( 0 ..size)
115+ . map ( |_| Arc :: new ( RwLock :: new ( SequenceCounter :: new ( ) ) ) )
116+ . collect ( )
117+ }
118+
53119 // transform add sequence nextval column.
54- async fn transform_sequence (
55- & self ,
120+ pub async fn transform_sequence (
121+ ctx : Arc < QueryContext > ,
56122 data_block : & mut DataBlock ,
123+ counter_lock : Arc < RwLock < SequenceCounter > > ,
57124 sequence_name : & String ,
58- data_type : & DataType ,
59125 ) -> Result < ( ) > {
60- transform_sequence ( & self . ctx , data_block, sequence_name, data_type) . await
126+ let count = data_block. num_rows ( ) as u64 ;
127+ let column = if count == 0 {
128+ UInt64Type :: from_data ( vec ! [ ] )
129+ } else {
130+ // Get or create the sequence counter
131+ let counter = counter_lock. read ( ) . await ;
132+
133+ // Try to reserve sequence numbers from the counter
134+ if let Some ( ( start, _end) ) = counter. try_reserve ( count) {
135+ // We have enough sequence numbers in the current batch
136+ let range = start..start + count;
137+ UInt64Type :: from_data ( range. collect :: < Vec < u64 > > ( ) )
138+ } else {
139+ // drop the read lock and get the write lock
140+ drop ( counter) ;
141+ let counter = counter_lock. write ( ) . await ;
142+ {
143+ // try reserve again
144+ if let Some ( ( start, _end) ) = counter. try_reserve ( count) {
145+ // We have enough sequence numbers in the current batch
146+ let range = start..start + count;
147+ UInt64Type :: from_data ( range. collect :: < Vec < u64 > > ( ) )
148+ } else {
149+ // We need to fetch more sequence numbers
150+ let tenant = ctx. get_tenant ( ) ;
151+ let catalog = ctx. get_default_catalog ( ) ?;
152+
153+ // Get current state of the counter
154+ let current = counter. current . load ( Ordering :: Relaxed ) ;
155+ let max = counter. max . load ( Ordering :: Relaxed ) ;
156+ // Calculate how many sequence numbers we need to fetch
157+ // If there are remaining numbers, we'll use them first
158+ let remaining = max. saturating_sub ( current) ;
159+ let to_fetch = count. saturating_sub ( remaining) ;
160+
161+ let step_size = ctx. get_settings ( ) . get_sequence_step_size ( ) ?;
162+ let batch_size = to_fetch. max ( step_size) ;
163+
164+ // Calculate batch size - take the larger of count or step_size
165+ let req = GetSequenceNextValueReq {
166+ ident : SequenceIdent :: new ( & tenant, sequence_name) ,
167+ count : batch_size,
168+ } ;
169+
170+ let resp = catalog. get_sequence_next_value ( req) . await ?;
171+ let start = resp. start ;
172+
173+ // If we have remaining numbers, use them first
174+ if remaining > 0 {
175+ // Then add the new batch after the remaining numbers
176+ counter. update_batch ( start, batch_size) ;
177+
178+ // Return a combined range: first the remaining numbers, then the new ones
179+ let mut numbers = Vec :: with_capacity ( count as usize ) ;
180+
181+ // Add the remaining numbers
182+ let remaining_to_use = remaining. min ( count) ;
183+ numbers. extend (
184+ ( current..current + remaining_to_use) . collect :: < Vec < u64 > > ( ) ,
185+ ) ;
186+
187+ // Add numbers from the new batch if needed
188+ if remaining_to_use < count {
189+ let new_needed = count - remaining_to_use;
190+ numbers. extend ( ( start..start + new_needed) . collect :: < Vec < u64 > > ( ) ) ;
191+ // Update the counter to reflect that we've used some of the new batch
192+ counter. current . store ( start + new_needed, Ordering :: SeqCst ) ;
193+ }
194+
195+ UInt64Type :: from_data ( numbers)
196+ } else {
197+ // No remaining numbers, just use the new batch
198+ counter. update_batch ( start + count, batch_size - count) ;
199+ // Return the sequence numbers needed for this request
200+ let range = start..start + count;
201+ UInt64Type :: from_data ( range. collect :: < Vec < u64 > > ( ) )
202+ }
203+ }
204+ }
205+ }
206+ } ;
207+
208+ data_block. add_column ( column) ;
209+ Ok ( ( ) )
61210 }
62211}
63212
@@ -70,10 +219,11 @@ impl AsyncTransform for TransformAsyncFunction {
70219 for ( i, async_func_desc) in self . async_func_descs . iter ( ) . enumerate ( ) {
71220 match & async_func_desc. func_arg {
72221 AsyncFunctionArgument :: SequenceFunction ( sequence_name) => {
73- self . transform_sequence (
222+ Self :: transform_sequence (
223+ self . ctx . clone ( ) ,
74224 & mut data_block,
225+ self . sequence_counters [ i] . clone ( ) ,
75226 sequence_name,
76- & async_func_desc. data_type ,
77227 )
78228 . await ?;
79229 }
@@ -92,28 +242,3 @@ impl AsyncTransform for TransformAsyncFunction {
92242 Ok ( data_block)
93243 }
94244}
95-
96- pub async fn transform_sequence (
97- ctx : & Arc < QueryContext > ,
98- data_block : & mut DataBlock ,
99- sequence_name : & String ,
100- _data_type : & DataType ,
101- ) -> Result < ( ) > {
102- let count = data_block. num_rows ( ) as u64 ;
103- let column = if count == 0 {
104- UInt64Type :: from_data ( vec ! [ ] )
105- } else {
106- let tenant = ctx. get_tenant ( ) ;
107- let catalog = ctx. get_default_catalog ( ) ?;
108- let req = GetSequenceNextValueReq {
109- ident : SequenceIdent :: new ( & tenant, sequence_name) ,
110- count,
111- } ;
112- let resp = catalog. get_sequence_next_value ( req) . await ?;
113- let range = resp. start ..resp. start + count;
114- UInt64Type :: from_data ( range. collect :: < Vec < u64 > > ( ) )
115- } ;
116- data_block. add_column ( column) ;
117-
118- Ok ( ( ) )
119- }
0 commit comments