common_query/
request.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
15mod base64_serde;
16mod initial_remote_dyn_filter_reg;
17
18use std::sync::Arc;
19
20use api::v1::region::RegionRequestHeader;
21use datafusion::execution::TaskContext;
22use datafusion::physical_expr::expressions::{Column, InListExpr, lit};
23use datafusion::physical_plan::PhysicalExpr;
24use datafusion::physical_plan::joins::HashTableLookupExpr;
25use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
26use datafusion_common::{DataFusionError, Result as DataFusionResult};
27use datafusion_expr::LogicalPlan;
28use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
29use datafusion_proto::physical_plan::from_proto::parse_physical_expr;
30use datafusion_proto::physical_plan::to_proto::serialize_physical_expr;
31use datafusion_proto::protobuf::PhysicalExprNode;
32use prost::Message;
33use serde::{Deserialize, Serialize};
34use snafu::ensure;
35use store_api::storage::RegionId;
36
37/// Current wire-format version for remote dynamic filter payload updates.
38pub use self::initial_remote_dyn_filter_reg::{
39    INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterReg,
40    InitialDynFilterRegs, InitialDynFilterSnapshot,
41};
42use crate::error::{DynFilterPayloadTooLargeSnafu, Error as CommonQueryError};
43
44pub const DYN_FILTER_PROTOCOL_VERSION: u32 = 1;
45/// Shared raw protobuf byte budget for remote dynamic filter payloads.
46pub const REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES: usize = 512 * 1024;
47
48/// Serialized predicate payload for remote dynamic filter updates.
49///
50/// The payload is tagged in JSON so receivers can reject unsupported encodings
51/// before decoding engine-specific bytes. For DataFusion expressions the
52/// `payload` bytes are serialized by `serde_json` as a base64 string, for example:
53///
54/// ```json
55/// { "kind": "datafusion", "payload": "CQgH" }
56/// ```
57#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
58#[non_exhaustive]
59#[serde(tag = "kind", content = "payload", rename_all = "snake_case")]
60pub enum DynFilterPayload {
61    /// A serialized DataFusion [`PhysicalExpr`] encoded as a protobuf
62    /// [`PhysicalExprNode`].
63    Datafusion(#[serde(with = "base64_serde::bytes")] Vec<u8>),
64}
65
66impl DynFilterPayload {
67    /// Encodes a DataFusion physical expression into a bounded dynamic filter payload.
68    ///
69    /// Runtime-only hash lookup predicates are degraded to `true` before encoding so
70    /// serializable min/max bounds around them can still be shipped to remote scans.
71    /// If the full serializable predicate is still larger than `max_payload_bytes`, large
72    /// membership predicates (`IN (...)`) are also degraded to `true` as a bounds-only fallback.
73    pub fn from_datafusion_expr(
74        expr: &Arc<dyn PhysicalExpr>,
75        max_payload_bytes: usize,
76    ) -> DataFusionResult<Self> {
77        match encode_remote_dyn_filter_expr(expr, max_payload_bytes, false) {
78            Ok(bytes) => Ok(Self::Datafusion(bytes)),
79            Err(CommonQueryError::DynFilterPayloadTooLarge { .. }) => {
80                encode_remote_dyn_filter_expr(expr, max_payload_bytes, true)
81                    .map(Self::Datafusion)
82                    .map_err(DataFusionError::from)
83            }
84            Err(error) => Err(DataFusionError::from(error)),
85        }
86    }
87
88    /// Decodes a DataFusion dynamic filter payload against the provided input schema.
89    ///
90    /// The decoded expression is validated to ensure column indexes stay within the receiver-side
91    /// schema, column names are consistent (defensive check), and the payload stays within
92    /// `max_payload_bytes`.
93    pub fn decode_datafusion_expr(
94        &self,
95        task_ctx: &TaskContext,
96        input_schema: &datafusion::arrow::datatypes::Schema,
97        max_payload_bytes: usize,
98    ) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
99        let Self::Datafusion(bytes) = self;
100        validate_payload_size(bytes.len(), max_payload_bytes).map_err(DataFusionError::from)?;
101        let codec = DefaultPhysicalExtensionCodec {};
102        let proto = PhysicalExprNode::decode(bytes.as_slice()).map_err(|e| {
103            DataFusionError::Internal(format!("Failed to decode PhysicalExprNode: {e}"))
104        })?;
105
106        let expr = parse_physical_expr(&proto, task_ctx, input_schema, &codec)?;
107        validate_supported_payload_expr(&expr)?;
108        validate_decoded_payload_expr(&expr, input_schema)?;
109        Ok(expr)
110    }
111}
112
113fn encode_physical_expr_to_bytes(expr: &Arc<dyn PhysicalExpr>) -> DataFusionResult<Vec<u8>> {
114    let codec = DefaultPhysicalExtensionCodec {};
115    let proto = serialize_physical_expr(expr, &codec)?;
116    let mut bytes = Vec::new();
117    proto.encode(&mut bytes).map_err(|e| {
118        DataFusionError::Internal(format!("Failed to encode PhysicalExprNode: {e}"))
119    })?;
120    Ok(bytes)
121}
122
123fn encode_remote_dyn_filter_expr(
124    expr: &Arc<dyn PhysicalExpr>,
125    max_payload_bytes: usize,
126    bounds_only: bool,
127) -> Result<Vec<u8>, CommonQueryError> {
128    let expr = portable_remote_dyn_filter_expr(Arc::clone(expr), bounds_only)
129        .map_err(CommonQueryError::from)?;
130    let bytes = encode_physical_expr_to_bytes(&expr).map_err(CommonQueryError::from)?;
131    validate_payload_size(bytes.len(), max_payload_bytes)?;
132    Ok(bytes)
133}
134
135fn portable_remote_dyn_filter_expr(
136    expr: Arc<dyn PhysicalExpr>,
137    bounds_only: bool,
138) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
139    expr.transform_up(|node| {
140        if node.is::<HashTableLookupExpr>() || (bounds_only && node.is::<InListExpr>()) {
141            Ok(Transformed::yes(lit(true)))
142        } else {
143            Ok(Transformed::no(node))
144        }
145    })
146    .map(|transformed| transformed.data)
147}
148
149pub(crate) fn decode_physical_expr_from_bytes(
150    bytes: &[u8],
151    task_ctx: &TaskContext,
152    input_schema: &datafusion::arrow::datatypes::Schema,
153    max_payload_bytes: usize,
154) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
155    validate_payload_size(bytes.len(), max_payload_bytes).map_err(DataFusionError::from)?;
156    let codec = DefaultPhysicalExtensionCodec {};
157    let proto = PhysicalExprNode::decode(bytes).map_err(|e| {
158        DataFusionError::Internal(format!("Failed to decode PhysicalExprNode: {e}"))
159    })?;
160
161    let expr = parse_physical_expr(&proto, task_ctx, input_schema, &codec)?;
162    validate_supported_payload_expr(&expr)?;
163    validate_decoded_payload_expr(&expr, input_schema)?;
164    Ok(expr)
165}
166
167fn validate_payload_size(
168    payload_size_bytes: usize,
169    max_payload_bytes: usize,
170) -> Result<(), CommonQueryError> {
171    ensure!(
172        payload_size_bytes <= max_payload_bytes,
173        DynFilterPayloadTooLargeSnafu {
174            payload_size_bytes,
175            max_payload_bytes,
176        }
177    );
178
179    Ok(())
180}
181
182fn validate_supported_payload_expr(expr: &Arc<dyn PhysicalExpr>) -> DataFusionResult<()> {
183    expr.apply(|node| {
184        if node.is::<HashTableLookupExpr>() {
185            return Err(DataFusionError::Plan(
186                "HashTableLookupExpr cannot be encoded into DynFilterPayload::Datafusion"
187                    .to_string(),
188            ));
189        }
190
191        Ok(TreeNodeRecursion::Continue)
192    })?;
193
194    Ok(())
195}
196
197/// Validates decoded dynamic filter physical expressions against the receiver-side schema.
198///
199/// Rejects out-of-bounds column indexes and, as a defensive check, rejects columns whose
200/// name disagrees with the corresponding schema field. DataFusion physical `Column` is
201/// index-authoritative, so a name mismatch usually indicates a coordinator/receiver
202/// schema inconsistency that should be surfaced loudly.
203fn validate_decoded_payload_expr(
204    expr: &Arc<dyn PhysicalExpr>,
205    input_schema: &datafusion::arrow::datatypes::Schema,
206) -> DataFusionResult<()> {
207    expr.apply(|node| {
208        if let Some(column) = node.downcast_ref::<Column>() {
209            let Some(field) = input_schema.fields().get(column.index()) else {
210                return Err(DataFusionError::Plan(format!(
211                    "Decoded Column '{}' references out-of-bounds index {} for input schema of size {}",
212                    column.name(),
213                    column.index(),
214                    input_schema.fields().len()
215                )));
216            };
217
218            if field.name() != column.name() {
219                return Err(DataFusionError::Plan(format!(
220                    "Decoded Column name/index mismatch: payload has '{}' at index {}, but schema field is '{}'",
221                    column.name(),
222                    column.index(),
223                    field.name()
224                )));
225            }
226        }
227
228        Ok(TreeNodeRecursion::Continue)
229    })?;
230
231    Ok(())
232}
233
234/// A remote dynamic filter update sent from a query coordinator to region servers.
235///
236/// `generation` is monotonic within a `query_id`/`filter_id` pair and matches the
237/// gRPC field name used by `RemoteDynFilterUpdate`. Receivers use it to ignore
238/// stale updates while `is_complete` marks the final payload for the filter.
239#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
240pub struct DynFilterUpdate {
241    /// Protocol version used by this update payload.
242    pub protocol_version: u32,
243    /// Internal query identifier that owns this dynamic filter lifecycle.
244    pub query_id: String,
245    /// Identifier of the dynamic filter within the query.
246    pub filter_id: String,
247    /// Monotonic update generation for this filter.
248    pub generation: u64,
249    /// Whether this update completes the dynamic filter stream.
250    pub is_complete: bool,
251    /// Serialized predicate payload carried by this update.
252    pub payload: DynFilterPayload,
253}
254
255impl DynFilterUpdate {
256    /// Creates a dynamic filter update with the current protocol version.
257    pub fn new(
258        query_id: String,
259        filter_id: String,
260        generation: u64,
261        is_complete: bool,
262        payload: DynFilterPayload,
263    ) -> Self {
264        Self {
265            protocol_version: DYN_FILTER_PROTOCOL_VERSION,
266            query_id,
267            filter_id,
268            generation,
269            is_complete,
270            payload,
271        }
272    }
273}
274
275/// The query request to be handled by the RegionServer (Datanode).
276#[derive(Clone, Debug)]
277pub struct QueryRequest {
278    /// The header of this request. Often to store some context of the query. None means all to defaults.
279    pub header: Option<RegionRequestHeader>,
280
281    /// The id of the region to be queried.
282    pub region_id: RegionId,
283
284    /// The form of the query: a logical plan.
285    pub plan: LogicalPlan,
286}
287
288#[cfg(test)]
289mod tests {
290    use std::sync::Arc;
291
292    use base64::Engine;
293    use base64::prelude::BASE64_STANDARD;
294    use datafusion::arrow::datatypes::{DataType, Field, Schema};
295    use datafusion::physical_expr::expressions::{BinaryExpr, Column, InListExpr, lit};
296    use datafusion::physical_plan::expressions::col;
297    use datafusion::physical_plan::joins::join_hash_map::JoinHashMapU32;
298    use datafusion::physical_plan::joins::{HashTableLookupExpr, Map, SeededRandomState};
299    use datafusion_expr::Operator;
300
301    use super::*;
302
303    #[test]
304    fn dyn_filter_update_sets_protocol_version() {
305        let update = DynFilterUpdate::new(
306            "query-1".to_string(),
307            "filter-1".to_string(),
308            3,
309            false,
310            DynFilterPayload::Datafusion(vec![1, 2, 3]),
311        );
312
313        assert_eq!(update.protocol_version, DYN_FILTER_PROTOCOL_VERSION);
314        assert!(!update.is_complete);
315        assert!(
316            matches!(update.payload, DynFilterPayload::Datafusion(ref bytes) if bytes == &vec![1, 2, 3])
317        );
318    }
319
320    #[test]
321    fn dyn_filter_update_json_round_trip_preserves_payload_shape() {
322        let update = DynFilterUpdate::new(
323            "query-2".to_string(),
324            "filter-9".to_string(),
325            9,
326            true,
327            DynFilterPayload::Datafusion(vec![9, 8, 7]),
328        );
329
330        let json = serde_json::to_string(&update).unwrap();
331        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
332        let decoded: DynFilterUpdate = serde_json::from_str(&json).unwrap();
333
334        assert_eq!(value["generation"], serde_json::json!(9));
335        assert!(value.get("epoch").is_none());
336        assert_eq!(
337            value["payload"],
338            serde_json::json!({ "kind": "datafusion", "payload": BASE64_STANDARD.encode([9, 8, 7]) })
339        );
340        assert_eq!(decoded, update);
341        assert!(decoded.is_complete);
342        assert!(
343            matches!(decoded.payload, DynFilterPayload::Datafusion(ref bytes) if bytes == &vec![9, 8, 7])
344        );
345    }
346
347    #[test]
348    fn dyn_filter_payload_json_uses_base64_for_empty_and_padded_payloads() {
349        let empty = serde_json::to_value(DynFilterPayload::Datafusion(vec![])).unwrap();
350        let one = serde_json::to_value(DynFilterPayload::Datafusion(vec![1])).unwrap();
351        let two = serde_json::to_value(DynFilterPayload::Datafusion(vec![1, 2])).unwrap();
352
353        assert_eq!(
354            empty,
355            serde_json::json!({"kind": "datafusion", "payload": ""})
356        );
357        assert_eq!(
358            one,
359            serde_json::json!({"kind": "datafusion", "payload": BASE64_STANDARD.encode([1])})
360        );
361        assert_eq!(
362            two,
363            serde_json::json!({"kind": "datafusion", "payload": BASE64_STANDARD.encode([1, 2])})
364        );
365    }
366
367    #[test]
368    fn dyn_filter_payload_json_rejects_invalid_base64() {
369        let err = serde_json::from_value::<DynFilterPayload>(serde_json::json!({
370            "kind": "datafusion",
371            "payload": "not base64!",
372        }))
373        .unwrap_err();
374
375        assert!(
376            err.to_string()
377                .contains("invalid base64 dynamic filter payload")
378        );
379    }
380
381    #[test]
382    fn dyn_filter_payload_round_trips_physical_column_expr() {
383        let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
384        let expr: Arc<dyn PhysicalExpr> =
385            Arc::new(Column::new_with_schema("host", &schema).unwrap());
386
387        let payload = DynFilterPayload::from_datafusion_expr(&expr, 1024).unwrap();
388        let decoded = payload
389            .decode_datafusion_expr(&TaskContext::default(), &schema, 1024)
390            .unwrap();
391
392        let original = expr.downcast_ref::<Column>().unwrap();
393        let decoded = decoded.downcast_ref::<Column>().unwrap();
394
395        assert_eq!(decoded.name(), original.name());
396        assert_eq!(decoded.index(), original.index());
397    }
398
399    #[test]
400    fn dyn_filter_payload_decode_rejects_invalid_bytes() {
401        let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
402        let payload = DynFilterPayload::Datafusion(vec![1, 2, 3]);
403
404        let err = payload
405            .decode_datafusion_expr(&TaskContext::default(), &schema, 1024)
406            .unwrap_err();
407
408        assert!(matches!(err, DataFusionError::Internal(_)));
409    }
410
411    #[test]
412    fn dyn_filter_payload_decode_rejects_column_name_index_mismatch() {
413        let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
414        let expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("service", 0));
415
416        let payload = DynFilterPayload::from_datafusion_expr(&expr, 1024).unwrap();
417        let err = payload
418            .decode_datafusion_expr(&TaskContext::default(), &schema, 1024)
419            .unwrap_err();
420
421        let msg = err.to_string();
422        assert!(
423            msg.contains("name/index mismatch"),
424            "expected name/index mismatch error, got: {msg}"
425        );
426        assert!(msg.contains("service"));
427        assert!(msg.contains("host"));
428    }
429
430    #[test]
431    fn dyn_filter_payload_decode_rejects_out_of_bounds_column_index() {
432        let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
433        let expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("host", 1));
434
435        let payload = DynFilterPayload::from_datafusion_expr(&expr, 1024).unwrap();
436        let err = payload
437            .decode_datafusion_expr(&TaskContext::default(), &schema, 1024)
438            .unwrap_err();
439
440        assert!(matches!(err, DataFusionError::Plan(_)));
441    }
442
443    #[test]
444    fn dyn_filter_payload_hash_lookup_fallback_preserves_bounds() {
445        let schema = Arc::new(Schema::new(vec![Field::new(
446            "device_id",
447            DataType::Int32,
448            false,
449        )]));
450        let device_id = col("device_id", &schema).unwrap();
451        let lower_bound = Arc::new(BinaryExpr::new(
452            Arc::clone(&device_id),
453            Operator::GtEq,
454            lit(10i32),
455        )) as Arc<dyn PhysicalExpr>;
456        let lookup = Arc::new(HashTableLookupExpr::new(
457            vec![Arc::clone(&device_id)],
458            SeededRandomState::with_seed(0),
459            Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))),
460            "hash_lookup".to_string(),
461        )) as Arc<dyn PhysicalExpr>;
462        let expr =
463            Arc::new(BinaryExpr::new(lower_bound, Operator::And, lookup)) as Arc<dyn PhysicalExpr>;
464
465        let payload = DynFilterPayload::from_datafusion_expr(&expr, 1024).unwrap();
466        let decoded = payload
467            .decode_datafusion_expr(&TaskContext::default(), &schema, 1024)
468            .unwrap();
469
470        assert!(!contains_expr::<HashTableLookupExpr>(&decoded));
471        let decoded_display = decoded.to_string();
472        assert!(decoded_display.contains("device_id"));
473        assert!(decoded_display.contains(">="));
474        assert!(!decoded_display.contains("hash_lookup"));
475    }
476
477    #[test]
478    fn dyn_filter_payload_oversized_inlist_falls_back_to_bounds() {
479        let schema = Arc::new(Schema::new(vec![Field::new(
480            "device_id",
481            DataType::Int32,
482            false,
483        )]));
484        let device_id = col("device_id", &schema).unwrap();
485        let lower_bound = Arc::new(BinaryExpr::new(
486            Arc::clone(&device_id),
487            Operator::GtEq,
488            lit(8192i32),
489        )) as Arc<dyn PhysicalExpr>;
490        let upper_bound = Arc::new(BinaryExpr::new(
491            Arc::clone(&device_id),
492            Operator::LtEq,
493            lit(8255i32),
494        )) as Arc<dyn PhysicalExpr>;
495        let bounds = Arc::new(BinaryExpr::new(lower_bound, Operator::And, upper_bound))
496            as Arc<dyn PhysicalExpr>;
497        let in_list = Arc::new(
498            InListExpr::try_new(
499                Arc::clone(&device_id),
500                (8192..8256).map(lit).collect(),
501                false,
502                &schema,
503            )
504            .unwrap(),
505        ) as Arc<dyn PhysicalExpr>;
506        let expr = Arc::new(BinaryExpr::new(Arc::clone(&bounds), Operator::And, in_list))
507            as Arc<dyn PhysicalExpr>;
508        let bounds_only = portable_remote_dyn_filter_expr(Arc::clone(&expr), true).unwrap();
509        let bounds_only_size = encode_physical_expr_to_bytes(&bounds_only).unwrap().len();
510        let full_size = encode_physical_expr_to_bytes(&expr).unwrap().len();
511        assert!(full_size > bounds_only_size);
512
513        let payload = DynFilterPayload::from_datafusion_expr(&expr, bounds_only_size).unwrap();
514        let decoded = payload
515            .decode_datafusion_expr(&TaskContext::default(), &schema, bounds_only_size)
516            .unwrap();
517
518        assert!(!contains_expr::<InListExpr>(&decoded));
519        let decoded_display = decoded.to_string();
520        assert!(decoded_display.contains("device_id"));
521        assert!(decoded_display.contains(">="));
522        assert!(decoded_display.contains("<="));
523    }
524
525    #[test]
526    fn dyn_filter_payload_rejects_oversized_payload() {
527        let expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("host", 0));
528
529        let err = DynFilterPayload::from_datafusion_expr(&expr, 1).unwrap_err();
530
531        let DataFusionError::External(error) = err else {
532            panic!("expected external common query error, got: {err:?}");
533        };
534        assert!(matches!(
535            error.downcast_ref::<CommonQueryError>(),
536            Some(CommonQueryError::DynFilterPayloadTooLarge { .. })
537        ));
538    }
539
540    fn contains_expr<T: PhysicalExpr>(expr: &Arc<dyn PhysicalExpr>) -> bool {
541        let mut found = false;
542        expr.apply(|node| {
543            if node.is::<T>() {
544                found = true;
545                Ok(TreeNodeRecursion::Stop)
546            } else {
547                Ok(TreeNodeRecursion::Continue)
548            }
549        })
550        .unwrap();
551        found
552    }
553}
common_query/request.rs

common_query/
request.rs
