feat: add remote dynamic filter frontend registration (#8148)

* feat: filter id

Signed-off-by: discord9 <discord9@163.com>

* feat: dyn filter registry

Signed-off-by: discord9 <discord9@163.com>

* feat: filter id&refactor to type

Signed-off-by: discord9 <discord9@163.com>

* feat: merge scan register dyn filter(not send yet)

Signed-off-by: discord9 <discord9@163.com>

* feat: init reg dyn filter

Signed-off-by: discord9 <discord9@163.com>

* wip: remote dyn filter task 03

Signed-off-by: discord9 <discord9@163.com>

* fix: resolve remote dyn filter rebase fallout

Signed-off-by: discord9 <discord9@163.com>

* chore: keep remote dyn filter docs local

Signed-off-by: discord9 <discord9@163.com>

* chore: remove stale filter id allow

Signed-off-by: discord9 <discord9@163.com>

* chore: clippy

Signed-off-by: discord9 <discord9@163.com>

* chore: fix remote dyn filter import style

Signed-off-by: discord9 <discord9@163.com>

* chore: fix query metrics test fallout

Signed-off-by: discord9 <discord9@163.com>

* fix: exclude region from remote dyn filter id

Signed-off-by: discord9 <discord9@163.com>

* chore: import

Signed-off-by: discord9 <discord9@163.com>

* refactor: rm some to latter

Signed-off-by: discord9 <discord9@163.com>

* feat: add initial dyn filter snapshot

Signed-off-by: discord9 <discord9@163.com>

* refactor: per review

Signed-off-by: discord9 <discord9@163.com>

* docs: better comment, rm some slop

Signed-off-by: discord9 <discord9@163.com>

* chore: per review

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-06-04 20:31:42 +08:00
committed by GitHub
parent 42eeeaa514
commit 6b7772e457
25 changed files with 2964 additions and 52 deletions
+1
View File
@@ -28,6 +28,7 @@ datatypes.workspace = true
once_cell.workspace = true
prost.workspace = true
serde.workspace = true
serde_json.workspace = true
snafu.workspace = true
sqlparser.workspace = true
sqlparser_derive = "0.1"
+40 -28
View File
@@ -12,10 +12,12 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod base64_serde;
mod initial_remote_dyn_filter_reg;
use std::sync::Arc;
use api::v1::region::RegionRequestHeader;
use datafusion::arrow::datatypes::Schema;
use datafusion::execution::TaskContext;
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_plan::PhysicalExpr;
@@ -32,6 +34,12 @@ use serde::{Deserialize, Serialize};
use store_api::storage::RegionId;
/// Current wire-format version for remote dynamic filter payload updates.
pub use self::initial_remote_dyn_filter_reg::{
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES, InitialDynFilterReg,
InitialDynFilterRegs, InitialDynFilterSnapshot,
};
pub const DYN_FILTER_PROTOCOL_VERSION: u32 = 1;
/// Serialized predicate payload for remote dynamic filter updates.
@@ -49,31 +57,7 @@ pub const DYN_FILTER_PROTOCOL_VERSION: u32 = 1;
pub enum DynFilterPayload {
/// A serialized DataFusion [`PhysicalExpr`] encoded as a protobuf
/// [`PhysicalExprNode`].
Datafusion(#[serde(with = "base64_bytes")] Vec<u8>),
}
mod base64_bytes {
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serializer};
pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&BASE64_STANDARD.encode(bytes))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let encoded = String::deserialize(deserializer)?;
BASE64_STANDARD.decode(encoded).map_err(|err| {
D::Error::custom(format!("invalid base64 dynamic filter payload: {err}"))
})
}
Datafusion(#[serde(with = "base64_serde::bytes")] Vec<u8>),
}
impl DynFilterPayload {
@@ -107,7 +91,7 @@ impl DynFilterPayload {
pub fn decode_datafusion_expr(
&self,
task_ctx: &TaskContext,
input_schema: &Schema,
input_schema: &datafusion::arrow::datatypes::Schema,
max_payload_bytes: usize,
) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
let Self::Datafusion(bytes) = self;
@@ -124,6 +108,34 @@ impl DynFilterPayload {
}
}
fn encode_physical_expr_to_bytes(expr: &Arc<dyn PhysicalExpr>) -> DataFusionResult<Vec<u8>> {
let codec = DefaultPhysicalExtensionCodec {};
let proto = serialize_physical_expr(expr, &codec)?;
let mut bytes = Vec::new();
proto.encode(&mut bytes).map_err(|e| {
DataFusionError::Internal(format!("Failed to encode PhysicalExprNode: {e}"))
})?;
Ok(bytes)
}
pub(crate) fn decode_physical_expr_from_bytes(
bytes: &[u8],
task_ctx: &TaskContext,
input_schema: &datafusion::arrow::datatypes::Schema,
max_payload_bytes: usize,
) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
validate_payload_size(bytes.len(), max_payload_bytes)?;
let codec = DefaultPhysicalExtensionCodec {};
let proto = PhysicalExprNode::decode(bytes).map_err(|e| {
DataFusionError::Internal(format!("Failed to decode PhysicalExprNode: {e}"))
})?;
let expr = parse_physical_expr(&proto, task_ctx, input_schema, &codec)?;
validate_supported_payload_expr(&expr)?;
validate_decoded_payload_expr(&expr, input_schema)?;
Ok(expr)
}
fn validate_payload_size(
payload_size_bytes: usize,
max_payload_bytes: usize,
@@ -161,7 +173,7 @@ fn validate_supported_payload_expr(expr: &Arc<dyn PhysicalExpr>) -> DataFusionRe
/// schema inconsistency that should be surfaced loudly.
fn validate_decoded_payload_expr(
expr: &Arc<dyn PhysicalExpr>,
input_schema: &Schema,
input_schema: &datafusion::arrow::datatypes::Schema,
) -> DataFusionResult<()> {
expr.apply(|node| {
if let Some(column) = node.as_any().downcast_ref::<Column>() {
@@ -0,0 +1,78 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! Serde adapters for byte fields encoded as base64 strings in JSON.
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
fn encode(bytes: &[u8]) -> String {
BASE64_STANDARD.encode(bytes)
}
fn decode(encoded: &str) -> Result<Vec<u8>, base64::DecodeError> {
BASE64_STANDARD.decode(encoded)
}
pub(crate) mod bytes {
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serializer};
pub fn serialize<S>(value: &[u8], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&super::encode(value))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let encoded = String::deserialize(deserializer)?;
super::decode(&encoded).map_err(|err| {
D::Error::custom(format!("invalid base64 dynamic filter payload: {err}"))
})
}
}
pub(crate) mod bytes_vec {
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub fn serialize<S>(values: &[Vec<u8>], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
values
.iter()
.map(|bytes| super::encode(bytes))
.collect::<Vec<_>>()
.serialize(serializer)
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<Vec<u8>>, D::Error>
where
D: Deserializer<'de>,
{
Vec::<String>::deserialize(deserializer)?
.into_iter()
.map(|encoded| {
super::decode(&encoded).map_err(|error| {
D::Error::custom(format!("invalid base64 bytes vector item: {error}"))
})
})
.collect()
}
}
@@ -0,0 +1,397 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::sync::Arc;
use datafusion::arrow::datatypes::Schema;
use datafusion::execution::TaskContext;
use datafusion::physical_plan::PhysicalExpr;
use datafusion_common::Result as DataFusionResult;
use serde::{Deserialize, Serialize};
use crate::request::{
DynFilterPayload, decode_physical_expr_from_bytes, encode_physical_expr_to_bytes,
};
pub const INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY: &str =
"initial_remote_dyn_filter_registrations";
pub const INITIAL_REMOTE_DYN_FILTER_REGS_MAX_COUNT: usize = 64;
/// Raw encoded registration byte budget for initial remote dynamic filter registrations.
///
/// Counts proto payload bytes before JSON/base64 expansion, not the final extension size.
pub const INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES: usize = 64 * 1024;
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct InitialDynFilterRegs {
#[serde(rename = "registrations")]
pub regs: Vec<InitialDynFilterReg>,
}
impl InitialDynFilterRegs {
pub fn new(regs: Vec<InitialDynFilterReg>) -> Self {
Self { regs }
}
pub fn is_empty(&self) -> bool {
self.regs.is_empty()
}
pub fn total_encoded_registration_bytes(&self) -> usize {
self.regs
.iter()
.map(InitialDynFilterReg::encoded_registration_bytes)
.sum()
}
pub fn validate_default_bounds(&self) -> Result<(), String> {
self.validate_bounds(
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_COUNT,
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES,
)
}
pub fn validate_bounds(
&self,
max_count: usize,
max_total_proto_bytes: usize,
) -> Result<(), String> {
if self.regs.len() > max_count {
return Err(format!(
"InitialDynFilterRegs contains {} registrations, which exceeds the configured limit of {}",
self.regs.len(),
max_count
));
}
let total_registration_bytes = self.total_encoded_registration_bytes();
if total_registration_bytes > max_total_proto_bytes {
return Err(format!(
"InitialDynFilterRegs contains {} total encoded registration bytes, which exceeds the configured limit of {}",
total_registration_bytes, max_total_proto_bytes
));
}
let mut seen_filter_ids = HashSet::with_capacity(self.regs.len());
for reg in &self.regs {
if !seen_filter_ids.insert(reg.filter_id.as_str()) {
return Err(format!(
"InitialDynFilterRegs contains duplicate filter_id '{}'",
reg.filter_id
));
}
}
Ok(())
}
pub fn to_extension_value(&self) -> serde_json::Result<String> {
serde_json::to_string(self)
}
pub fn from_extension_value(value: &str) -> serde_json::Result<Self> {
let regs = serde_json::from_str::<Self>(value)?;
regs.validate_default_bounds()
.map_err(serde::de::Error::custom)?;
Ok(regs)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct InitialDynFilterReg {
pub filter_id: String,
#[serde(with = "super::base64_serde::bytes_vec")]
pub child_exprs_datafusion_proto: Vec<Vec<u8>>,
/// Optional producer-side predicate snapshot captured at initial registration time.
///
/// This is only an initial pending update for the remote runtime filter. It is not part of
/// registration identity; identity is carried by `filter_id` and child expressions.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub initial_snapshot: Option<InitialDynFilterSnapshot>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct InitialDynFilterSnapshot {
pub payload: DynFilterPayload,
/// Producer-side generation used to ignore stale snapshots.
pub generation: u64,
/// Whether this snapshot completes the dynamic filter stream.
pub is_complete: bool,
}
impl InitialDynFilterReg {
pub fn new(filter_id: impl Into<String>, child_exprs_datafusion_proto: Vec<Vec<u8>>) -> Self {
Self {
filter_id: filter_id.into(),
child_exprs_datafusion_proto,
initial_snapshot: None,
}
}
pub fn with_initial_snapshot(mut self, initial_snapshot: InitialDynFilterSnapshot) -> Self {
self.initial_snapshot = Some(initial_snapshot);
self
}
pub fn from_filter_id_and_children(
filter_id: impl Into<String>,
children: &[Arc<dyn PhysicalExpr>],
) -> DataFusionResult<Self> {
let child_exprs_datafusion_proto = children
.iter()
.map(encode_physical_expr_to_bytes)
.collect::<DataFusionResult<Vec<_>>>()?;
Ok(Self::new(filter_id, child_exprs_datafusion_proto))
}
pub fn encoded_child_expr_bytes(&self) -> usize {
self.child_exprs_datafusion_proto.iter().map(Vec::len).sum()
}
pub fn encoded_registration_bytes(&self) -> usize {
self.encoded_child_expr_bytes()
+ self
.initial_snapshot
.as_ref()
.map(InitialDynFilterSnapshot::encoded_payload_bytes)
.unwrap_or(0)
}
pub fn decode_children(
&self,
task_ctx: &TaskContext,
input_schema: &Schema,
max_payload_bytes: usize,
) -> DataFusionResult<Vec<Arc<dyn PhysicalExpr>>> {
self.child_exprs_datafusion_proto
.iter()
.map(|expr_bytes| {
decode_physical_expr_from_bytes(
expr_bytes,
task_ctx,
input_schema,
max_payload_bytes,
)
})
.collect::<DataFusionResult<Vec<_>>>()
}
}
impl InitialDynFilterSnapshot {
pub fn new(payload: DynFilterPayload, generation: u64, is_complete: bool) -> Self {
Self {
payload,
generation,
is_complete,
}
}
pub fn encoded_payload_bytes(&self) -> usize {
match &self.payload {
DynFilterPayload::Datafusion(bytes) => bytes.len(),
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_plan::PhysicalExpr;
use datafusion_common::DataFusionError;
use super::*;
#[test]
fn initial_dyn_filter_regs_json_round_trip() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3]]),
InitialDynFilterReg::new("filter-b", vec![vec![4, 5]]),
]);
let encoded = regs.to_extension_value().unwrap();
let json: serde_json::Value = serde_json::from_str(&encoded).unwrap();
let decoded = InitialDynFilterRegs::from_extension_value(&encoded).unwrap();
assert_eq!(
json["registrations"][0]["child_exprs_datafusion_proto"],
serde_json::json!(["AQID"])
);
assert_eq!(
json["registrations"][1]["child_exprs_datafusion_proto"],
serde_json::json!(["BAU="])
);
assert_eq!(decoded, regs);
}
#[test]
fn initial_dyn_filter_regs_json_round_trip_with_snapshot() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3]]).with_initial_snapshot(
InitialDynFilterSnapshot::new(DynFilterPayload::Datafusion(vec![4, 5, 6]), 7, true),
),
]);
let encoded = regs.to_extension_value().unwrap();
let json: serde_json::Value = serde_json::from_str(&encoded).unwrap();
let decoded = InitialDynFilterRegs::from_extension_value(&encoded).unwrap();
assert_eq!(
json["registrations"][0]["child_exprs_datafusion_proto"],
serde_json::json!(["AQID"])
);
assert_eq!(
json["registrations"][0]["initial_snapshot"]["payload"],
serde_json::json!({"kind":"datafusion","payload":"BAUG"})
);
assert_eq!(decoded, regs);
assert_eq!(
decoded.regs[0]
.initial_snapshot
.as_ref()
.unwrap()
.generation,
7
);
assert!(
decoded.regs[0]
.initial_snapshot
.as_ref()
.unwrap()
.is_complete
);
}
#[test]
fn initial_dyn_filter_reg_json_defaults_missing_snapshot_to_none() {
let decoded = InitialDynFilterRegs::from_extension_value(
r#"{"registrations":[{"filter_id":"filter-a","child_exprs_datafusion_proto":["AQID"]}]}"#,
)
.unwrap();
assert_eq!(decoded.regs.len(), 1);
assert!(decoded.regs[0].initial_snapshot.is_none());
}
#[test]
fn initial_dyn_filter_reg_encoded_registration_bytes_include_snapshot_payload() {
let reg = InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3], vec![4]])
.with_initial_snapshot(InitialDynFilterSnapshot::new(
DynFilterPayload::Datafusion(vec![5, 6]),
2,
false,
));
assert_eq!(reg.encoded_child_expr_bytes(), 4);
assert_eq!(reg.encoded_registration_bytes(), 6);
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_duplicate_filter_ids() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1]]),
InitialDynFilterReg::new("filter-a", vec![vec![2]]),
]);
let err = regs.validate_bounds(8, 1024).unwrap_err();
assert!(err.contains("duplicate filter_id 'filter-a'"));
}
#[test]
fn initial_dyn_filter_regs_from_extension_value_validates_default_bounds() {
let value = r#"{"registrations":[{"filter_id":"filter-a","child_exprs_datafusion_proto":["AQ=="]},{"filter_id":"filter-a","child_exprs_datafusion_proto":["Ag=="]}]}"#;
let err = InitialDynFilterRegs::from_extension_value(value).unwrap_err();
assert!(err.to_string().contains("duplicate filter_id 'filter-a'"));
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_too_many_regs() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1]]),
InitialDynFilterReg::new("filter-b", vec![vec![2]]),
]);
let err = regs.validate_bounds(1, 1024).unwrap_err();
assert!(err.contains("exceeds the configured limit of 1"));
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_total_proto_bytes_over_limit() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3]]),
InitialDynFilterReg::new("filter-b", vec![vec![4, 5, 6]]),
]);
let err = regs.validate_bounds(8, 5).unwrap_err();
assert!(err.contains("6 total encoded registration bytes"));
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_snapshot_bytes_over_limit() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1]]).with_initial_snapshot(
InitialDynFilterSnapshot::new(
DynFilterPayload::Datafusion(vec![2, 3, 4]),
2,
false,
),
),
]);
let err = regs.validate_bounds(8, 3).unwrap_err();
assert!(err.contains("4 total encoded registration bytes"));
}
#[test]
fn initial_dyn_filter_reg_round_trips_child_exprs() {
let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
let child: Arc<dyn PhysicalExpr> =
Arc::new(Column::new_with_schema("host", &schema).unwrap());
let reg = InitialDynFilterReg::from_filter_id_and_children("filter-1", &[child]).unwrap();
let decoded = reg
.decode_children(&TaskContext::default(), &schema, 1024)
.unwrap();
let decoded = decoded[0].as_any().downcast_ref::<Column>().unwrap();
assert_eq!(reg.filter_id, "filter-1");
assert_eq!(decoded.name(), "host");
assert_eq!(decoded.index(), 0);
}
#[test]
fn initial_dyn_filter_reg_decode_rejects_column_name_index_mismatch() {
let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
let reg = InitialDynFilterReg::from_filter_id_and_children(
"filter-1",
&[Arc::new(Column::new("service", 0)) as Arc<dyn PhysicalExpr>],
)
.unwrap();
let err = reg
.decode_children(&TaskContext::default(), &schema, 1024)
.unwrap_err();
assert!(matches!(err, DataFusionError::Plan(_)));
}
}