datanode/region_server/catalog/
inspect.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Streaming providers for internal inspection tables.
16
17use std::sync::Arc;
18
19use datafusion::arrow::datatypes::SchemaRef;
20use datafusion::catalog::{Session, TableProvider};
21use datafusion::error::Result;
22use datafusion::execution::TaskContext;
23use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
24use datafusion::physical_plan::streaming::{PartitionStream, StreamingTableExec};
25use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
26use datafusion_common::DataFusionError;
27use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType};
28use futures::StreamExt;
29use mito2::engine::MitoEngine;
30use store_api::compaction_task::{CompactionTaskBatchBuilder, CompactionTaskEntry};
31use store_api::sst_entry::{ManifestSstBatchBuilder, ManifestSstEntry};
32
33/// Supplies batches for an internal inspection table. Implementations may prune
34/// conservatively; DataFusion retains complete residual predicates.
35pub(crate) trait InspectSource: std::fmt::Debug + Send + Sync {
36    fn schema(&self) -> SchemaRef;
37
38    /// Creates a fresh scan at execution time, including on repeated execution.
39    /// Projection includes columns required by residual predicates.
40    fn scan(&self, request: &InspectScan, ctx: Arc<TaskContext>) -> SendableRecordBatchStream;
41}
42
43/// Parameters shared by inspection sources. Limits are enforced by the physical
44/// plan, after residual filters, rather than by individual sources.
45#[derive(Debug, Clone)]
46pub(crate) struct InspectScan {
47    schema: SchemaRef,
48    projection: Option<Vec<usize>>,
49    filters: Vec<Expr>,
50    batch_size: usize,
51}
52
53/// Adapts a lazy inspection source to DataFusion's streaming execution plan.
54#[derive(Debug)]
55pub(crate) struct InspectTable {
56    source: Arc<dyn InspectSource>,
57}
58
59impl InspectTable {
60    /// Wraps a source without opening a scan.
61    pub(crate) fn new(source: Arc<dyn InspectSource>) -> Self {
62        Self { source }
63    }
64}
65
66#[async_trait::async_trait]
67impl TableProvider for InspectTable {
68    fn schema(&self) -> SchemaRef {
69        self.source.schema()
70    }
71
72    fn table_type(&self) -> TableType {
73        TableType::View
74    }
75
76    fn supports_filters_pushdown(
77        &self,
78        filters: &[&Expr],
79    ) -> Result<Vec<TableProviderFilterPushDown>> {
80        Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
81    }
82
83    async fn scan(
84        &self,
85        state: &dyn Session,
86        projection: Option<&Vec<usize>>,
87        filters: &[Expr],
88        limit: Option<usize>,
89    ) -> Result<Arc<dyn ExecutionPlan>> {
90        // An inexact predicate must be checked before selecting a limited prefix.
91        let limit = limit.filter(|_| filters.is_empty());
92        let configured_batch_size: usize = state.config_options().execution.batch_size.into();
93        let schema = match projection {
94            Some(projection) => Arc::new(self.schema().project(projection)?),
95            None => self.schema(),
96        };
97        let partition = InspectPartition {
98            source: self.source.clone(),
99            request: InspectScan {
100                schema: schema.clone(),
101                projection: projection.cloned(),
102                filters: filters.to_vec(),
103                batch_size: limit
104                    .unwrap_or(configured_batch_size)
105                    .min(configured_batch_size)
106                    .max(1),
107            },
108        };
109        Ok(Arc::new(StreamingTableExec::try_new(
110            schema,
111            vec![Arc::new(partition)],
112            None,
113            None,
114            false,
115            limit,
116        )?))
117    }
118}
119
120/// Reads manifest snapshots without materializing the full SST inventory.
121pub(crate) struct ManifestSstSource {
122    engine: MitoEngine,
123}
124
125impl ManifestSstSource {
126    /// Creates a source backed by the engine's open-region manifests.
127    pub(crate) fn new(engine: MitoEngine) -> Self {
128        Self { engine }
129    }
130}
131
132impl std::fmt::Debug for ManifestSstSource {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        f.debug_struct("ManifestSstSource").finish_non_exhaustive()
135    }
136}
137
138impl InspectSource for ManifestSstSource {
139    fn schema(&self) -> SchemaRef {
140        ManifestSstEntry::schema().arrow_schema().clone()
141    }
142
143    fn scan(&self, request: &InspectScan, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
144        let builder = ManifestSstBatchBuilder::new(request.schema.clone());
145        let stream = self
146            .engine
147            .scan_ssts_from_manifest_projected(
148                &request.filters,
149                request.batch_size,
150                request.projection.as_deref(),
151            )
152            .map(move |entries| builder.build(&entries).map_err(DataFusionError::from));
153        Box::pin(RecordBatchStreamAdapter::new(
154            request.schema.clone(),
155            stream,
156        ))
157    }
158}
159
160/// Reads live compaction state only when a scan is executed.
161pub(crate) struct CompactionTasksSource {
162    engine: MitoEngine,
163}
164
165impl CompactionTasksSource {
166    /// Creates a source backed by the engine's live compaction cycles.
167    pub(crate) fn new(engine: MitoEngine) -> Self {
168        Self { engine }
169    }
170}
171
172impl std::fmt::Debug for CompactionTasksSource {
173    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        f.debug_struct("CompactionTasksSource")
175            .finish_non_exhaustive()
176    }
177}
178
179impl InspectSource for CompactionTasksSource {
180    fn schema(&self) -> SchemaRef {
181        CompactionTaskEntry::schema().arrow_schema().clone()
182    }
183
184    fn scan(&self, request: &InspectScan, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
185        let engine = self.engine.clone();
186        let request = request.clone();
187        let schema = request.schema.clone();
188        let builder = CompactionTaskBatchBuilder::new(schema.clone());
189        let stream = async_stream::try_stream! {
190            let entries = engine
191                .inspect_compaction_tasks(&request.filters)
192                .await
193                .map_err(|error| DataFusionError::External(Box::new(error)))?;
194            for entries in entries.chunks(request.batch_size) {
195                yield builder.build(entries)?;
196            }
197        };
198        Box::pin(RecordBatchStreamAdapter::new(schema, stream))
199    }
200}
201
202#[derive(Debug)]
203struct InspectPartition {
204    source: Arc<dyn InspectSource>,
205    request: InspectScan,
206}
207
208impl PartitionStream for InspectPartition {
209    fn schema(&self) -> &SchemaRef {
210        &self.request.schema
211    }
212
213    fn execute(&self, ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
214        self.source.scan(&self.request, ctx)
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use std::sync::atomic::{AtomicUsize, Ordering};
221
222    use datafusion::arrow::array::{ArrayRef, Int32Array};
223    use datafusion::arrow::datatypes::{DataType, Field, Schema};
224    use datafusion::arrow::record_batch::RecordBatch;
225    use datafusion::prelude::SessionContext;
226
227    use super::*;
228
229    #[derive(Debug, Default)]
230    struct CountingSource {
231        scans: Arc<AtomicUsize>,
232        rows: Arc<AtomicUsize>,
233        drops: Arc<AtomicUsize>,
234    }
235
236    struct ScanGuard(Arc<AtomicUsize>);
237
238    impl Drop for ScanGuard {
239        fn drop(&mut self) {
240            self.0.fetch_add(1, Ordering::SeqCst);
241        }
242    }
243
244    impl InspectSource for CountingSource {
245        fn schema(&self) -> SchemaRef {
246            Arc::new(Schema::new(vec![
247                Field::new("id", DataType::Int32, false),
248                Field::new("value", DataType::Int32, false),
249            ]))
250        }
251
252        fn scan(&self, request: &InspectScan, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
253            self.scans.fetch_add(1, Ordering::SeqCst);
254            let guard = ScanGuard(self.drops.clone());
255            let rows = self.rows.clone();
256            let schema = request.schema.clone();
257            let projection = request.projection.clone().unwrap_or_else(|| vec![0, 1]);
258            // Deliberately ignore source predicates: the provider must preserve
259            // residual filtering before limit, including hidden predicate columns.
260            let stream = futures::stream::iter(0..100).map(move |id| {
261                let _guard = &guard;
262                rows.fetch_add(1, Ordering::SeqCst);
263                let columns: Vec<ArrayRef> = projection
264                    .iter()
265                    .map(|column| {
266                        Arc::new(Int32Array::from(vec![if *column == 0 {
267                            id
268                        } else {
269                            id * 10
270                        }])) as ArrayRef
271                    })
272                    .collect();
273                RecordBatch::try_new(schema.clone(), columns).map_err(DataFusionError::from)
274            });
275            Box::pin(RecordBatchStreamAdapter::new(
276                request.schema.clone(),
277                stream,
278            ))
279        }
280    }
281
282    #[tokio::test]
283    async fn test_lazy_repeatable_scan_and_limit_cancellation() {
284        let source = Arc::new(CountingSource::default());
285        let ctx = SessionContext::new();
286        ctx.register_table("inspection", Arc::new(InspectTable::new(source.clone())))
287            .unwrap();
288        let frame = ctx
289            .sql("SELECT value FROM inspection LIMIT 2")
290            .await
291            .unwrap();
292        let plan = frame.create_physical_plan().await.unwrap();
293        assert_eq!(source.scans.load(Ordering::SeqCst), 0);
294        assert_eq!(source.rows.load(Ordering::SeqCst), 0);
295        for run in 1..=2 {
296            let batches = datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx())
297                .await
298                .unwrap();
299            let values = batches
300                .iter()
301                .flat_map(|batch| {
302                    batch
303                        .column(0)
304                        .as_any()
305                        .downcast_ref::<Int32Array>()
306                        .unwrap()
307                        .values()
308                        .to_vec()
309                })
310                .collect::<Vec<_>>();
311            assert_eq!(values, vec![0, 10]);
312            assert_eq!(source.scans.load(Ordering::SeqCst), run);
313            assert_eq!(source.rows.load(Ordering::SeqCst), run * 2);
314            assert_eq!(source.drops.load(Ordering::SeqCst), run);
315        }
316    }
317
318    #[tokio::test]
319    async fn test_hidden_residual_filter_before_limit() {
320        let source = Arc::new(CountingSource::default());
321        let ctx = SessionContext::new();
322        ctx.register_table("inspection", Arc::new(InspectTable::new(source.clone())))
323            .unwrap();
324        let batches = ctx
325            .sql("SELECT value FROM inspection WHERE id >= 8 LIMIT 2")
326            .await
327            .unwrap()
328            .collect()
329            .await
330            .unwrap();
331        let values = batches
332            .iter()
333            .flat_map(|batch| {
334                assert_eq!(batch.num_columns(), 1);
335                batch
336                    .column(0)
337                    .as_any()
338                    .downcast_ref::<Int32Array>()
339                    .unwrap()
340                    .values()
341                    .to_vec()
342            })
343            .collect::<Vec<_>>();
344        // No ORDER BY: repartitioning may choose any two qualifying rows.
345        assert_eq!(values.len(), 2);
346        assert!(
347            values
348                .iter()
349                .all(|value| (80..1000).contains(value) && value % 10 == 0)
350        );
351        assert_eq!(source.scans.load(Ordering::SeqCst), 1);
352        assert_eq!(source.drops.load(Ordering::SeqCst), 1);
353    }
354    #[tokio::test]
355    async fn test_manifest_optimized_sql_pruning() {
356        use std::sync::Mutex;
357
358        use datafusion_expr::{col, lit};
359        use mito2::config::MitoConfig;
360        use mito2::manifest::action::RegionEdit;
361        use mito2::sst::file::FileMeta;
362        use mito2::test_util::{CreateRequestBuilder, TestEnv};
363        use store_api::region_engine::RegionEngine;
364        use store_api::region_request::RegionRequest;
365        use store_api::storage::{FileId, RegionId};
366
367        #[derive(Debug)]
368        struct RecordingSource {
369            inner: ManifestSstSource,
370            filters: Mutex<Vec<Expr>>,
371            candidates: Arc<AtomicUsize>,
372        }
373
374        impl InspectSource for RecordingSource {
375            fn schema(&self) -> SchemaRef {
376                self.inner.schema()
377            }
378
379            fn scan(
380                &self,
381                request: &InspectScan,
382                ctx: Arc<TaskContext>,
383            ) -> SendableRecordBatchStream {
384                *self.filters.lock().unwrap() = request.filters.clone();
385                let candidates = self.candidates.clone();
386                let stream = self.inner.scan(request, ctx).map(move |batch| {
387                    if let Ok(batch) = &batch {
388                        candidates.fetch_add(batch.num_rows(), Ordering::SeqCst);
389                    }
390                    batch
391                });
392                Box::pin(RecordBatchStreamAdapter::new(
393                    request.schema.clone(),
394                    stream,
395                ))
396            }
397        }
398
399        let mut env = TestEnv::new().await.with_node_id(42);
400        let engine = env
401            .create_engine(MitoConfig {
402                schedule_compaction_after_edit: false,
403                ..Default::default()
404            })
405            .await;
406        let region_id = RegionId::new(11, 1);
407        engine
408            .handle_request(
409                region_id,
410                RegionRequest::Create(CreateRequestBuilder::new().build()),
411            )
412            .await
413            .unwrap();
414        engine
415            .edit_region(
416                region_id,
417                RegionEdit {
418                    files_to_add: vec![FileMeta {
419                        region_id,
420                        file_id: FileId::from_bytes(1u128.to_be_bytes()),
421                        ..Default::default()
422                    }],
423                    files_to_remove: vec![],
424                    timestamp_ms: None,
425                    compaction_time_window: None,
426                    flushed_entry_id: None,
427                    flushed_sequence: None,
428                    committed_sequence: None,
429                    wal_recovery: Default::default(),
430                },
431            )
432            .await
433            .unwrap();
434        let source = Arc::new(RecordingSource {
435            inner: ManifestSstSource::new(engine),
436            filters: Mutex::new(Vec::new()),
437            candidates: Arc::new(AtomicUsize::new(0)),
438        });
439        let ctx = SessionContext::new();
440        ctx.register_table("manifest", Arc::new(InspectTable::new(source.clone())))
441            .unwrap();
442        for (predicate, expected_filter, expected_count) in [
443            ("node_id = 99".to_string(), col("node_id").eq(lit(99u64)), 0),
444            (
445                "node_id != 42".to_string(),
446                col("node_id").not_eq(lit(42u64)),
447                0,
448            ),
449            ("node_id = 42".to_string(), col("node_id").eq(lit(42u64)), 1),
450            (
451                "region_id = 0".to_string(),
452                col("region_id").eq(lit(0u64)),
453                0,
454            ),
455            (
456                format!("region_id = {}", region_id.as_u64()),
457                col("region_id").eq(lit(region_id.as_u64())),
458                1,
459            ),
460        ] {
461            source.candidates.store(0, Ordering::SeqCst);
462            source.filters.lock().unwrap().clear();
463            // collect() runs the default optimizer and the actual inspection provider.
464            let batches = ctx
465                .sql(&format!("SELECT file_id FROM manifest WHERE {predicate}"))
466                .await
467                .unwrap()
468                .collect()
469                .await
470                .unwrap();
471            // The scanner receives a UInt64 comparison, with no Decimal cast.
472            assert_eq!(
473                *source.filters.lock().unwrap(),
474                vec![expected_filter],
475                "{predicate}"
476            );
477            // Count the manifest source's candidates before DataFusion's residual filter.
478            assert_eq!(
479                source.candidates.load(Ordering::SeqCst),
480                expected_count,
481                "{predicate}"
482            );
483            assert_eq!(
484                batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
485                expected_count,
486                "{predicate}"
487            );
488        }
489    }
490}
datanode/region_server/catalog/inspect.rs

datanode/region_server/catalog/
inspect.rs
