store_api/
compaction_task.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
15use std::sync::Arc;
16
17use common_recordbatch::DfRecordBatch;
18use common_recordbatch::filter::MetadataFilter;
19use common_time::Timestamp;
20use common_time::timestamp::TimeUnit;
21use datafusion_common::{DataFusionError, ScalarValue};
22use datafusion_expr::LogicalPlan;
23use datatypes::arrow::array::{
24    ArrayRef, StringArray, TimestampMillisecondArray, TimestampNanosecondArray, UInt32Array,
25    UInt64Array,
26};
27use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
28use datatypes::arrow::error::ArrowError;
29use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
30
31use crate::inspect::{self, ProjectedBatchBuilder};
32use crate::storage::{RegionId, ScanRequest, TableId};
33
34// Compaction task column names.
35pub const TASK_ID: &str = "task_id";
36pub const REGION_ID: &str = "region_id";
37pub const TABLE_ID: &str = "table_id";
38pub const NODE_ID: &str = "node_id";
39pub const TRIGGER_TYPE: &str = "trigger_type";
40pub const STATE: &str = "state";
41pub const PHASE: &str = "phase";
42pub const SUBMITTED_AT: &str = "submitted_at";
43pub const ELAPSED_MILLIS: &str = "elapsed_millis";
44pub const INPUT_FILES: &str = "input_files";
45pub const INPUT_BYTES: &str = "input_bytes";
46pub const TOTAL_UNITS: &str = "total_units";
47pub const COMPLETED_UNITS: &str = "completed_units";
48pub const REQUEST_OPTIONS: &str = "request_options";
49pub const EFFECTIVE_OPTIONS: &str = "effective_options";
50pub const INPUT_TIME_START: &str = "input_time_start";
51pub const INPUT_TIME_END: &str = "input_time_end";
52pub const OUTPUT_TIME_START: &str = "output_time_start";
53pub const OUTPUT_TIME_END: &str = "output_time_end";
54
55/// Queued compaction cycle.
56pub const STATE_QUEUED: &str = "queued";
57/// Running compaction cycle.
58pub const STATE_RUNNING: &str = "running";
59/// Compaction cycle with cancellation requested.
60pub const STATE_CANCELLING: &str = "cancelling";
61/// Cycle waiting for execution.
62pub const PHASE_WAITING: &str = "waiting";
63/// Cycle selecting a compaction plan.
64pub const PHASE_PLANNING: &str = "planning";
65/// Cycle executing local compaction units.
66pub const PHASE_EXECUTING: &str = "executing";
67/// Cycle committing compaction output.
68pub const PHASE_COMMITTING: &str = "committing";
69/// Cycle executing through a remote compactor.
70pub const PHASE_REMOTE: &str = "remote";
71/// Compaction requested manually.
72pub const TRIGGER_TYPE_MANUAL: &str = "manual";
73/// Compaction triggered automatically.
74pub const TRIGGER_TYPE_AUTOMATIC: &str = "automatic";
75
76/// A live region-level compaction cycle. Completed cycles are not retained.
77///
78/// String fields are shared so cloning a snapshot does not copy parameter payloads.
79/// Time ranges are envelopes and may contain holes. SQL endpoints are NULL when
80/// unknown, unbounded, or outside the representable nanosecond timestamp range.
81#[derive(Debug, Clone)]
82pub struct CompactionTaskEntry {
83    /// UUID retained from acceptance until this cycle leaves the scheduler.
84    pub task_id: Arc<str>,
85    /// Physical region whose scheduler owns this cycle.
86    pub region_id: RegionId,
87    /// Physical table ID derived from the region ID.
88    pub table_id: TableId,
89    /// Reporting datanode ID (0 in standalone), not the remote executor ID.
90    pub node_id: Option<u64>,
91    /// Whether this cycle was triggered automatically or manually.
92    pub trigger_type: Arc<str>,
93    /// Queued, running, or cancelling, aggregated across the cycle.
94    pub state: Arc<str>,
95    /// Phase label: `waiting`, `planning`, `executing`, `committing`, or `remote`.
96    pub phase: Arc<str>,
97    /// Acceptance time in milliseconds since the Unix epoch.
98    pub submitted_at: i64,
99    /// Monotonic elapsed time since acceptance, including waiting.
100    pub elapsed_millis: u64,
101    /// Distinct plan input files, including expired SSTs; unknown before plan acceptance.
102    pub input_files: Option<u64>,
103    /// Total size of plan inputs, including expired SSTs; not bytes processed.
104    pub input_bytes: Option<u64>,
105    /// Number of local atomic units; unavailable for remote execution.
106    pub total_units: Option<u64>,
107    /// Local units with a terminal result, including failures.
108    pub completed_units: Option<u64>,
109    /// JSON snapshot of request arguments, including requested parallelism.
110    pub request_options: Arc<str>,
111    /// JSON snapshot of resolved settings and picker output, if available.
112    pub effective_options: Option<Arc<str>>,
113    /// Inclusive earliest timestamp in plan input SSTs, including expired SSTs.
114    pub input_time_start: Option<Timestamp>,
115    /// Inclusive latest timestamp in plan input SSTs, including expired SSTs.
116    pub input_time_end: Option<Timestamp>,
117    /// Inclusive start of the planned output envelope, not actual output SST bounds.
118    pub output_time_start: Option<Timestamp>,
119    /// Exclusive end of the planned output envelope, not actual output SST bounds.
120    pub output_time_end: Option<Timestamp>,
121}
122
123impl CompactionTaskEntry {
124    /// Returns the public inspection schema.
125    pub fn schema() -> SchemaRef {
126        use datatypes::prelude::ConcreteDataType as Ty;
127        Arc::new(Schema::new(vec![
128            ColumnSchema::new(TASK_ID, Ty::string_datatype(), false),
129            ColumnSchema::new(REGION_ID, Ty::uint64_datatype(), false),
130            ColumnSchema::new(TABLE_ID, Ty::uint32_datatype(), false),
131            ColumnSchema::new(NODE_ID, Ty::uint64_datatype(), true),
132            ColumnSchema::new(TRIGGER_TYPE, Ty::string_datatype(), false),
133            ColumnSchema::new(STATE, Ty::string_datatype(), false),
134            ColumnSchema::new(PHASE, Ty::string_datatype(), false),
135            ColumnSchema::new(SUBMITTED_AT, Ty::timestamp_millisecond_datatype(), false),
136            ColumnSchema::new(ELAPSED_MILLIS, Ty::uint64_datatype(), false),
137            ColumnSchema::new(INPUT_FILES, Ty::uint64_datatype(), true),
138            ColumnSchema::new(INPUT_BYTES, Ty::uint64_datatype(), true),
139            ColumnSchema::new(TOTAL_UNITS, Ty::uint64_datatype(), true),
140            ColumnSchema::new(COMPLETED_UNITS, Ty::uint64_datatype(), true),
141            ColumnSchema::new(REQUEST_OPTIONS, Ty::string_datatype(), false),
142            ColumnSchema::new(EFFECTIVE_OPTIONS, Ty::string_datatype(), true),
143            ColumnSchema::new(INPUT_TIME_START, Ty::timestamp_nanosecond_datatype(), true),
144            ColumnSchema::new(INPUT_TIME_END, Ty::timestamp_nanosecond_datatype(), true),
145            ColumnSchema::new(OUTPUT_TIME_START, Ty::timestamp_nanosecond_datatype(), true),
146            ColumnSchema::new(OUTPUT_TIME_END, Ty::timestamp_nanosecond_datatype(), true),
147        ]))
148    }
149
150    /// Prunes using known values, leaving unsupported predicates to DataFusion.
151    pub fn matches(&self, filters: &MetadataFilter) -> bool {
152        filters.matches(|name| {
153            Some(match name {
154                TASK_ID => ScalarValue::Utf8(Some(self.task_id.to_string())),
155                REGION_ID => ScalarValue::UInt64(Some(self.region_id.as_u64())),
156                TABLE_ID => ScalarValue::UInt32(Some(self.table_id)),
157                NODE_ID => ScalarValue::UInt64(self.node_id),
158                TRIGGER_TYPE => ScalarValue::Utf8(Some(self.trigger_type.to_string())),
159                STATE => ScalarValue::Utf8(Some(self.state.to_string())),
160                PHASE => ScalarValue::Utf8(Some(self.phase.to_string())),
161                SUBMITTED_AT => ScalarValue::TimestampMillisecond(Some(self.submitted_at), None),
162                ELAPSED_MILLIS => ScalarValue::UInt64(Some(self.elapsed_millis)),
163                INPUT_FILES => ScalarValue::UInt64(self.input_files),
164                INPUT_BYTES => ScalarValue::UInt64(self.input_bytes),
165                TOTAL_UNITS => ScalarValue::UInt64(self.total_units),
166                COMPLETED_UNITS => ScalarValue::UInt64(self.completed_units),
167                REQUEST_OPTIONS => ScalarValue::Utf8(Some(self.request_options.to_string())),
168                EFFECTIVE_OPTIONS => {
169                    ScalarValue::Utf8(self.effective_options.as_deref().map(str::to_string))
170                }
171                INPUT_TIME_START => {
172                    ScalarValue::TimestampNanosecond(timestamp_nanos(self.input_time_start), None)
173                }
174                INPUT_TIME_END => {
175                    ScalarValue::TimestampNanosecond(timestamp_nanos(self.input_time_end), None)
176                }
177                OUTPUT_TIME_START => {
178                    ScalarValue::TimestampNanosecond(timestamp_nanos(self.output_time_start), None)
179                }
180                OUTPUT_TIME_END => {
181                    ScalarValue::TimestampNanosecond(timestamp_nanos(self.output_time_end), None)
182                }
183                _ => return None,
184            })
185        })
186    }
187
188    /// Converts compaction task entries to a record batch.
189    pub fn to_record_batch(entries: &[Self]) -> Result<DfRecordBatch, ArrowError> {
190        Self::to_projected_record_batch(entries, None)
191    }
192
193    /// Builds only the requested columns, including empty and repeated projections.
194    fn to_projected_record_batch(
195        entries: &[Self],
196        projection: Option<&[usize]>,
197    ) -> Result<DfRecordBatch, ArrowError> {
198        let schema = Self::schema().arrow_schema().clone();
199        let schema = match projection {
200            Some(projection) => Arc::new(schema.project(projection)?),
201            None => schema,
202        };
203        CompactionTaskBatchBuilder::new(schema).build(entries)
204    }
205
206    /// Internal name, inaccessible through user table resolution.
207    pub fn reserved_table_name_for_inspection() -> &'static str {
208        "__inspect/__mito/__compaction_tasks"
209    }
210
211    /// Builds a remote scan with filtering before projection and limit.
212    pub fn build_plan(request: ScanRequest) -> Result<LogicalPlan, DataFusionError> {
213        inspect::build_plan(
214            request,
215            Self::reserved_table_name_for_inspection(),
216            Self::schema(),
217        )
218    }
219}
220
221/// Converts compaction task batches using a schema prepared once per scan.
222/// Repeated projected columns share the same Arrow array.
223pub struct CompactionTaskBatchBuilder {
224    projected: ProjectedBatchBuilder,
225}
226
227impl CompactionTaskBatchBuilder {
228    /// Creates a builder for a projected compaction task schema.
229    pub fn new(schema: ArrowSchemaRef) -> Self {
230        Self {
231            projected: ProjectedBatchBuilder::new(schema),
232        }
233    }
234
235    /// Builds the projected columns, copying string contents only into Arrow buffers.
236    pub fn build(&self, entries: &[CompactionTaskEntry]) -> Result<DfRecordBatch, ArrowError> {
237        self.projected.build(entries.len(), |name| {
238            let array: ArrayRef = match name {
239                TASK_ID => Arc::new(StringArray::from_iter_values(
240                    entries.iter().map(|e| e.task_id.as_ref()),
241                )),
242                REGION_ID => Arc::new(UInt64Array::from_iter_values(
243                    entries.iter().map(|e| e.region_id.as_u64()),
244                )),
245                TABLE_ID => Arc::new(UInt32Array::from_iter_values(
246                    entries.iter().map(|e| e.table_id),
247                )),
248                NODE_ID => Arc::new(UInt64Array::from_iter(entries.iter().map(|e| e.node_id))),
249                TRIGGER_TYPE => Arc::new(StringArray::from_iter_values(
250                    entries.iter().map(|e| e.trigger_type.as_ref()),
251                )),
252                STATE => Arc::new(StringArray::from_iter_values(
253                    entries.iter().map(|e| e.state.as_ref()),
254                )),
255                PHASE => Arc::new(StringArray::from_iter_values(
256                    entries.iter().map(|e| e.phase.as_ref()),
257                )),
258                SUBMITTED_AT => Arc::new(TimestampMillisecondArray::from_iter_values(
259                    entries.iter().map(|e| e.submitted_at),
260                )),
261                ELAPSED_MILLIS => Arc::new(UInt64Array::from_iter_values(
262                    entries.iter().map(|e| e.elapsed_millis),
263                )),
264                INPUT_FILES => Arc::new(UInt64Array::from_iter(
265                    entries.iter().map(|e| e.input_files),
266                )),
267                INPUT_BYTES => Arc::new(UInt64Array::from_iter(
268                    entries.iter().map(|e| e.input_bytes),
269                )),
270                TOTAL_UNITS => Arc::new(UInt64Array::from_iter(
271                    entries.iter().map(|e| e.total_units),
272                )),
273                COMPLETED_UNITS => Arc::new(UInt64Array::from_iter(
274                    entries.iter().map(|e| e.completed_units),
275                )),
276                REQUEST_OPTIONS => Arc::new(StringArray::from_iter_values(
277                    entries.iter().map(|e| e.request_options.as_ref()),
278                )),
279                EFFECTIVE_OPTIONS => Arc::new(StringArray::from_iter(
280                    entries.iter().map(|e| e.effective_options.as_deref()),
281                )),
282                INPUT_TIME_START => Arc::new(TimestampNanosecondArray::from_iter(
283                    entries.iter().map(|e| timestamp_nanos(e.input_time_start)),
284                )),
285                INPUT_TIME_END => Arc::new(TimestampNanosecondArray::from_iter(
286                    entries.iter().map(|e| timestamp_nanos(e.input_time_end)),
287                )),
288                OUTPUT_TIME_START => Arc::new(TimestampNanosecondArray::from_iter(
289                    entries.iter().map(|e| timestamp_nanos(e.output_time_start)),
290                )),
291                OUTPUT_TIME_END => Arc::new(TimestampNanosecondArray::from_iter(
292                    entries.iter().map(|e| timestamp_nanos(e.output_time_end)),
293                )),
294                _ => {
295                    return Err(ArrowError::SchemaError(format!(
296                        "Unknown compaction task column: {name}"
297                    )));
298                }
299            };
300            Ok(array)
301        })
302    }
303}
304
305/// Converts endpoints identically for predicate pruning and Arrow output.
306fn timestamp_nanos(timestamp: Option<Timestamp>) -> Option<i64> {
307    timestamp
308        .and_then(|ts| ts.convert_to(TimeUnit::Nanosecond))
309        .map(|ts| ts.value())
310}
311
312#[cfg(test)]
313mod tests {
314    use datafusion_expr::{col, lit};
315    use datatypes::arrow::datatypes::DataType;
316
317    use super::*;
318
319    fn entry() -> CompactionTaskEntry {
320        CompactionTaskEntry {
321            task_id: "task-1".into(),
322            region_id: RegionId::new(1, 42),
323            table_id: 1,
324            node_id: Some(7),
325            trigger_type: "manual".into(),
326            state: "queued".into(),
327            phase: "waiting".into(),
328            submitted_at: 123,
329            elapsed_millis: 10,
330            input_files: Some(3),
331            input_bytes: Some(4096),
332            total_units: Some(5),
333            completed_units: Some(2),
334            request_options: "{\"request\":1}".into(),
335            effective_options: Some("{\"effective\":2}".into()),
336            input_time_start: Some(Timestamp::new_nanosecond(11)),
337            input_time_end: Some(Timestamp::new_nanosecond(22)),
338            output_time_start: Some(Timestamp::new_nanosecond(33)),
339            output_time_end: Some(Timestamp::new_nanosecond(44)),
340        }
341    }
342
343    #[test]
344    fn test_compaction_task_clone_shares_strings() {
345        let entry = entry();
346        let cloned = entry.clone();
347        for (original, cloned) in [
348            (&entry.task_id, &cloned.task_id),
349            (&entry.trigger_type, &cloned.trigger_type),
350            (&entry.state, &cloned.state),
351            (&entry.phase, &cloned.phase),
352            (&entry.request_options, &cloned.request_options),
353            (
354                entry.effective_options.as_ref().unwrap(),
355                cloned.effective_options.as_ref().unwrap(),
356            ),
357        ] {
358            assert!(Arc::ptr_eq(original, cloned));
359        }
360        assert_eq!(
361            CompactionTaskEntry::to_record_batch(&[entry]).unwrap(),
362            CompactionTaskEntry::to_record_batch(&[cloned]).unwrap()
363        );
364    }
365
366    #[test]
367    fn test_compaction_task_column_values_and_predicates() {
368        let entry = entry();
369        let batch = CompactionTaskEntry::to_record_batch(std::slice::from_ref(&entry)).unwrap();
370        // Independent SQL contract values distinguish every same-typed field.
371        let expected = [
372            ("task_id", ScalarValue::Utf8(Some("task-1".into()))),
373            ("region_id", ScalarValue::UInt64(Some(4_294_967_338))),
374            ("table_id", ScalarValue::UInt32(Some(1))),
375            ("node_id", ScalarValue::UInt64(Some(7))),
376            ("trigger_type", ScalarValue::Utf8(Some("manual".into()))),
377            ("state", ScalarValue::Utf8(Some("queued".into()))),
378            ("phase", ScalarValue::Utf8(Some("waiting".into()))),
379            (
380                "submitted_at",
381                ScalarValue::TimestampMillisecond(Some(123), None),
382            ),
383            ("elapsed_millis", ScalarValue::UInt64(Some(10))),
384            ("input_files", ScalarValue::UInt64(Some(3))),
385            ("input_bytes", ScalarValue::UInt64(Some(4096))),
386            ("total_units", ScalarValue::UInt64(Some(5))),
387            ("completed_units", ScalarValue::UInt64(Some(2))),
388            (
389                "request_options",
390                ScalarValue::Utf8(Some("{\"request\":1}".into())),
391            ),
392            (
393                "effective_options",
394                ScalarValue::Utf8(Some("{\"effective\":2}".into())),
395            ),
396            (
397                "input_time_start",
398                ScalarValue::TimestampNanosecond(Some(11), None),
399            ),
400            (
401                "input_time_end",
402                ScalarValue::TimestampNanosecond(Some(22), None),
403            ),
404            (
405                "output_time_start",
406                ScalarValue::TimestampNanosecond(Some(33), None),
407            ),
408            (
409                "output_time_end",
410                ScalarValue::TimestampNanosecond(Some(44), None),
411            ),
412        ];
413        assert_eq!(batch.num_columns(), expected.len());
414        for (index, (name, value)) in expected.into_iter().enumerate() {
415            assert_eq!(batch.schema().field(index).name(), name);
416            assert_eq!(
417                ScalarValue::try_from_array(batch.column(index), 0).unwrap(),
418                value,
419                "{name}"
420            );
421            assert!(
422                entry.matches(&MetadataFilter::new(&[col(name).eq(lit(value.clone()))])),
423                "{name}"
424            );
425            let different = match &value {
426                ScalarValue::Utf8(Some(value)) => ScalarValue::Utf8(Some(format!("{value}-other"))),
427                ScalarValue::UInt32(Some(value)) => ScalarValue::UInt32(Some(value + 1)),
428                ScalarValue::UInt64(Some(value)) => ScalarValue::UInt64(Some(value + 1)),
429                ScalarValue::TimestampMillisecond(Some(value), timezone) => {
430                    ScalarValue::TimestampMillisecond(Some(value + 1), timezone.clone())
431                }
432                ScalarValue::TimestampNanosecond(Some(value), timezone) => {
433                    ScalarValue::TimestampNanosecond(Some(value + 1), timezone.clone())
434                }
435                unexpected => panic!("unexpected contract value for {name}: {unexpected:?}"),
436            };
437            assert!(
438                !entry.matches(&MetadataFilter::new(&[col(name).eq(lit(different))])),
439                "{name}"
440            );
441            assert!(
442                !entry.matches(&MetadataFilter::new(&[col(name).not_eq(lit(value))])),
443                "{name}"
444            );
445        }
446    }
447
448    #[test]
449    fn test_compaction_task_schema_and_projection() {
450        let entries = [entry()];
451        let full = CompactionTaskEntry::to_record_batch(&entries).unwrap();
452        let fields: Vec<_> = full
453            .schema()
454            .fields()
455            .iter()
456            .map(|f| (f.name().clone(), f.data_type().clone(), f.is_nullable()))
457            .collect();
458        let expected = [
459            ("task_id", DataType::Utf8, false),
460            ("region_id", DataType::UInt64, false),
461            ("table_id", DataType::UInt32, false),
462            ("node_id", DataType::UInt64, true),
463            ("trigger_type", DataType::Utf8, false),
464            ("state", DataType::Utf8, false),
465            ("phase", DataType::Utf8, false),
466            (
467                "submitted_at",
468                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None),
469                false,
470            ),
471            ("elapsed_millis", DataType::UInt64, false),
472            ("input_files", DataType::UInt64, true),
473            ("input_bytes", DataType::UInt64, true),
474            ("total_units", DataType::UInt64, true),
475            ("completed_units", DataType::UInt64, true),
476            ("request_options", DataType::Utf8, false),
477            ("effective_options", DataType::Utf8, true),
478            (
479                "input_time_start",
480                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Nanosecond, None),
481                true,
482            ),
483            (
484                "input_time_end",
485                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Nanosecond, None),
486                true,
487            ),
488            (
489                "output_time_start",
490                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Nanosecond, None),
491                true,
492            ),
493            (
494                "output_time_end",
495                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Nanosecond, None),
496                true,
497            ),
498        ]
499        .map(|(name, ty, nullable)| (name.to_string(), ty, nullable));
500        assert_eq!(fields, expected);
501        for projection in [
502            vec![],
503            vec![5, 0, 5],
504            vec![14, 3, 8],
505            vec![18, 15, 17, 16, 15],
506        ] {
507            let actual =
508                CompactionTaskEntry::to_projected_record_batch(&entries, Some(&projection))
509                    .unwrap();
510            assert_eq!(actual, full.project(&projection).unwrap());
511            assert_eq!(actual.num_rows(), 1);
512        }
513        assert!(CompactionTaskEntry::to_projected_record_batch(&entries, Some(&[99])).is_err());
514        let projected_schema = Arc::new(full.schema().project(&[0, 9, 0, 18, 9]).unwrap());
515        let builder = CompactionTaskBatchBuilder::new(projected_schema.clone());
516        for rows in [entries.as_slice(), &entries[..0]] {
517            let projected = builder.build(rows).unwrap();
518            assert_eq!(
519                projected,
520                CompactionTaskEntry::to_record_batch(rows)
521                    .unwrap()
522                    .project(&[0, 9, 0, 18, 9])
523                    .unwrap()
524            );
525            assert!(Arc::ptr_eq(&projected.schema(), &projected_schema));
526            assert!(Arc::ptr_eq(projected.column(0), projected.column(2)));
527            assert!(Arc::ptr_eq(projected.column(1), projected.column(4)));
528        }
529    }
530
531    #[test]
532    fn test_compaction_task_time_ranges_precision_and_nulls() {
533        let precise = CompactionTaskEntry {
534            input_time_start: Some(Timestamp::new(-2, TimeUnit::Second)),
535            input_time_end: Some(Timestamp::new(3, TimeUnit::Millisecond)),
536            output_time_start: Some(Timestamp::new(-4, TimeUnit::Microsecond)),
537            output_time_end: Some(Timestamp::new(5, TimeUnit::Nanosecond)),
538            ..entry()
539        };
540        let overflow = CompactionTaskEntry {
541            input_time_start: Some(Timestamp::new(i64::MIN, TimeUnit::Second)),
542            input_time_end: Some(Timestamp::new(i64::MAX, TimeUnit::Second)),
543            output_time_start: Some(Timestamp::new(i64::MIN, TimeUnit::Millisecond)),
544            output_time_end: Some(Timestamp::new(i64::MAX, TimeUnit::Microsecond)),
545            ..entry()
546        };
547        let unknown = CompactionTaskEntry {
548            input_time_start: None,
549            input_time_end: None,
550            output_time_start: None,
551            output_time_end: None,
552            ..entry()
553        };
554        let entries = [precise, overflow, unknown];
555        let batch =
556            CompactionTaskEntry::to_projected_record_batch(&entries, Some(&[15, 16, 17, 18]))
557                .unwrap();
558        for (index, (name, nanos)) in [
559            ("input_time_start", -2_000_000_000),
560            ("input_time_end", 3_000_000),
561            ("output_time_start", -4_000),
562            ("output_time_end", 5),
563        ]
564        .into_iter()
565        .enumerate()
566        {
567            let values = batch
568                .column(index)
569                .as_any()
570                .downcast_ref::<TimestampNanosecondArray>()
571                .unwrap();
572            assert_eq!(
573                values,
574                &TimestampNanosecondArray::from(vec![Some(nanos), None, None])
575            );
576            let equal = col(name).eq(lit(ScalarValue::TimestampNanosecond(Some(nanos), None)));
577            let different =
578                col(name).eq(lit(ScalarValue::TimestampNanosecond(Some(nanos + 1), None)));
579            assert!(entries[0].matches(&MetadataFilter::new(std::slice::from_ref(&equal))));
580            assert!(!entries[0].matches(&MetadataFilter::new(&[different])));
581            assert!(!entries[0].matches(&MetadataFilter::new(&[col(name).is_null()])));
582            for entry in &entries[1..] {
583                assert!(entry.matches(&MetadataFilter::new(&[col(name).is_null()])));
584                assert!(!entry.matches(&MetadataFilter::new(&[col(name).is_not_null()])));
585                assert!(!entry.matches(&MetadataFilter::new(std::slice::from_ref(&equal))));
586            }
587        }
588        let full = CompactionTaskEntry::to_record_batch(&entries).unwrap();
589        let projection = [18, 16, 15, 18, 17];
590        assert_eq!(
591            CompactionTaskEntry::to_projected_record_batch(&entries, Some(&projection)).unwrap(),
592            full.project(&projection).unwrap()
593        );
594    }
595
596    #[test]
597    fn test_compaction_task_predicates_and_plan() {
598        let entry = entry();
599        for (expr, expected) in [
600            (col("task_id").eq(lit("task-1")), true),
601            (col("table_id").eq(lit(2u32)), false),
602            (col("state").eq(lit("running")), false),
603            (col("node_id").is_null(), false),
604            (
605                col("state")
606                    .eq(lit("running"))
607                    .or(col("unknown").eq(lit(1))),
608                true,
609            ),
610        ] {
611            assert_eq!(entry.matches(&MetadataFilter::new(&[expr])), expected);
612        }
613        let plan = CompactionTaskEntry::build_plan(ScanRequest {
614            projection: Some(vec![0]),
615            filters: vec![col("state").eq(lit("running"))],
616            limit: Some(1),
617            ..Default::default()
618        })
619        .unwrap();
620        assert_eq!(plan.schema().fields().len(), 1);
621        let plan = plan.display_indent().to_string();
622        assert!(plan.contains("Filter:"), "{plan}");
623        assert!(plan.contains("Limit:"), "{plan}");
624    }
625}
store_api/compaction_task.rs

store_api/
compaction_task.rs
