common_function/admin/
discard_unflushed_data.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 arrow::datatypes::DataType as ArrowDataType;
16use common_error::ext::BoxedError;
17use common_macro::admin_fn;
18use common_query::error::{
19    InvalidFuncArgsSnafu, MissingTableMutationHandlerSnafu, Result, TableMutationSnafu,
20    UnsupportedInputDataTypeSnafu,
21};
22use datafusion_expr::{Signature, TypeSignature, Volatility};
23use datatypes::data_type::DataType;
24use datatypes::prelude::*;
25use session::context::QueryContextRef;
26use session::table_name::table_name_to_full_name;
27use snafu::{ResultExt, ensure};
28use store_api::storage::RegionId;
29use table::table_name::TableName;
30
31use crate::handlers::TableMutationHandlerRef;
32use crate::helper::cast_u64;
33
34/// Discards all unflushed data from a region.
35#[admin_fn(
36    name = DiscardUnflushedDataFunction,
37    display_name = discard_unflushed,
38    sig_fn = signature,
39    ret = uint64,
40    single_row
41)]
42pub(crate) async fn discard_unflushed_data(
43    table_mutation_handler: &TableMutationHandlerRef,
44    query_ctx: &QueryContextRef,
45    params: &[ValueRef<'_>],
46) -> Result<Value> {
47    ensure!(
48        params.len() == 1,
49        InvalidFuncArgsSnafu {
50            err_msg: format!(
51                "The length of the args is not correct, expect 1, have: {}",
52                params.len()
53            ),
54        }
55    );
56
57    let affected_rows = match params[0] {
58        ValueRef::String(table_name) => {
59            let (catalog_name, schema_name, table_name) =
60                table_name_to_full_name(table_name, query_ctx)
61                    .map_err(BoxedError::new)
62                    .context(TableMutationSnafu)?;
63            table_mutation_handler
64                .discard_unflushed_data_by_table(
65                    TableName::new(catalog_name, schema_name, table_name),
66                    query_ctx.clone(),
67                )
68                .await?
69        }
70        _ => {
71            let Some(region_id) = cast_u64(&params[0])? else {
72                return UnsupportedInputDataTypeSnafu {
73                    function: "discard_unflushed",
74                    datatypes: params
75                        .iter()
76                        .map(|value| value.data_type())
77                        .collect::<Vec<_>>(),
78                }
79                .fail();
80            };
81            table_mutation_handler
82                .discard_unflushed_data(RegionId::from_u64(region_id), query_ctx.clone())
83                .await?
84        }
85    };
86
87    Ok(Value::from(affected_rows as u64))
88}
89
90fn signature() -> Signature {
91    Signature::one_of(
92        vec![
93            TypeSignature::Uniform(
94                1,
95                ConcreteDataType::numerics()
96                    .into_iter()
97                    .map(|data_type| data_type.as_arrow_type())
98                    .collect(),
99            ),
100            TypeSignature::Exact(vec![ArrowDataType::Utf8]),
101        ],
102        Volatility::Immutable,
103    )
104}
105
106#[cfg(test)]
107mod tests {
108    use std::sync::Arc;
109
110    use arrow::array::{StringArray, UInt64Array};
111    use arrow::datatypes::{DataType, Field};
112    use datafusion_expr::{ColumnarValue, TypeSignature};
113
114    use super::*;
115    use crate::function::FunctionContext;
116    use crate::function_factory::ScalarFunctionFactory;
117    use crate::function_registry::{FUNCTION_REGISTRY, get_admin_function};
118
119    #[test]
120    fn test_discard_unflushed_data_is_admin_only() {
121        assert!(get_admin_function("discard_unflushed").is_some());
122        assert!(
123            FUNCTION_REGISTRY
124                .get_function("discard_unflushed")
125                .is_none()
126        );
127    }
128
129    #[test]
130    fn test_discard_unflushed_data_signature() {
131        let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into();
132        let function = factory.provide(FunctionContext::mock());
133
134        assert_eq!("discard_unflushed", function.name());
135        assert_eq!(DataType::UInt64, function.return_type(&[]).unwrap());
136        assert!(matches!(
137            function.signature(),
138            Signature {
139                type_signature: TypeSignature::OneOf(valid_types),
140                volatility: Volatility::Immutable,
141                ..
142            } if valid_types == &vec![
143                TypeSignature::Uniform(
144                    1,
145                    ConcreteDataType::numerics()
146                        .into_iter()
147                        .map(|data_type| {
148                            use datatypes::data_type::DataType;
149                            data_type.as_arrow_type()
150                        })
151                        .collect::<Vec<_>>(),
152                ),
153                TypeSignature::Exact(vec![DataType::Utf8]),
154            ]
155        ));
156    }
157
158    #[tokio::test]
159    async fn test_discard_unflushed_data() {
160        let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into();
161        let function = factory.provide(FunctionContext::mock());
162        let args = datafusion::logical_expr::ScalarFunctionArgs {
163            args: vec![ColumnarValue::Array(Arc::new(UInt64Array::from(vec![99])))],
164            arg_fields: vec![Arc::new(Field::new("arg_0", DataType::UInt64, false))],
165            return_field: Arc::new(Field::new("result", DataType::UInt64, false)),
166            number_rows: 1,
167            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
168        };
169
170        let result = function
171            .as_async()
172            .unwrap()
173            .invoke_async_with_args(args)
174            .await
175            .unwrap();
176        let ColumnarValue::Array(array) = result else {
177            panic!("expected array output");
178        };
179        let array = array.as_any().downcast_ref::<UInt64Array>().unwrap();
180        assert_eq!(42, array.value(0));
181    }
182
183    #[tokio::test]
184    async fn test_discard_unflushed_data_by_table() {
185        let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into();
186        let function = factory.provide(FunctionContext::mock());
187        let args = datafusion::logical_expr::ScalarFunctionArgs {
188            args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
189                "my_table",
190            ])))],
191            arg_fields: vec![Arc::new(Field::new("arg_0", DataType::Utf8, false))],
192            return_field: Arc::new(Field::new("result", DataType::UInt64, false)),
193            number_rows: 1,
194            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
195        };
196
197        let result = function
198            .as_async()
199            .unwrap()
200            .invoke_async_with_args(args)
201            .await
202            .unwrap();
203        let ColumnarValue::Array(array) = result else {
204            panic!("expected array output");
205        };
206        let array = array.as_any().downcast_ref::<UInt64Array>().unwrap();
207        assert_eq!(42, array.value(0));
208    }
209
210    #[tokio::test]
211    async fn test_discard_unflushed_data_rejects_multiple_rows() {
212        let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into();
213        let function = factory.provide(FunctionContext::mock());
214        let args = datafusion::logical_expr::ScalarFunctionArgs {
215            args: vec![ColumnarValue::Array(Arc::new(UInt64Array::from(vec![
216                1, 2,
217            ])))],
218            arg_fields: vec![Arc::new(Field::new("arg_0", DataType::UInt64, false))],
219            return_field: Arc::new(Field::new("result", DataType::UInt64, false)),
220            number_rows: 2,
221            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
222        };
223
224        let error = function
225            .as_async()
226            .unwrap()
227            .invoke_async_with_args(args)
228            .await
229            .unwrap_err();
230        assert!(error.to_string().contains("received 2"));
231    }
232}
common_function/admin/discard_unflushed_data.rs

common_function/admin/
discard_unflushed_data.rs
