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

common_query/
request.rs
