starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-04 16:09:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: upgrade datafusion to 54 (#12665)

Browse Source 
- [x] ~add sort when to many groups, we can do this, but for now do not
do this~
- [x] remove the unsafe for memory pool
- [x] add approx_distinct improve, already fork, will add in next pr
- [x] fork vortex to add cast fix, and export
`.exclude_schemes([IntDictScheme.id()])`
- [x] ~enable_physical_uncorrelated_scalar_subquery test~
This commit is contained in:
Huaijin
2026-06-16 13:27:26 +00:00
committed by GitHub
parent 7207ea1aba
commit 0831badd02
60 changed files with 715 additions and 753 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+318 -245
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+5 -5
View File
@@ -319,10 +319,10 @@ cityhasher = { version = "0.1", default-features = false }
collapse = "0.1.2"
cron = "0.15"
dashmap = { version = "6.1", features = ["serde"] }
datafusion = "53"
datafusion-proto = "53"
datafusion-functions-aggregate-common = "53"
datafusion-functions-json = "0.53"
datafusion = "54"
datafusion-proto = "54"
datafusion-functions-aggregate-common = "54"
datafusion-functions-json = "0.54"
expect-test = "1.4"
arrow = { version = "58", features = ["ipc_compression", "prettyprint"] }
arrow-flight = "58"
@@ -330,7 +330,7 @@ arrow-json = "58"
arrow-schema = { version = "58", features = ["serde"] }
parquet = { version = "58", features = ["arrow", "async", "object_store"] }
object_store = { version = "0.13.2", features = ["aws", "azure", "gcp"] }
sqlparser = { version = "0.61", features = ["serde", "visitor"] }
sqlparser = { version = "0.62", features = ["serde", "visitor"] }
dotenv_config = "0.2"
dotenvy = "0.15"
env_logger = "0.11"
src/config/Cargo.toml
+1 -1
View File
@@ -90,7 +90,7 @@ vrl.workspace = true
tonic.workspace = true
tonic-prost.workspace = true
tracing-opentelemetry.workspace = true
vortex = { git = "https://github.com/openobserve/vortex", rev = "13423ca204f81dc1b20ca040da2aeb386220a312", features = [
vortex = { git = "https://github.com/openobserve/vortex", branch = "0.75.0-fix", features = [
"tokio",
], optional = true }
src/config/src/utils/file.rs
+3 -1
View File
@@ -431,7 +431,9 @@ mod tests {
assert!(result.is_err());
// Test range where start > end (this should not cause overflow)
let result = get_file_contents(file_path.to_str().unwrap(), Some(5..3));
let start = 5;
let end = 3;
let result = get_file_contents(file_path.to_str().unwrap(), Some(start..end));
assert!(result.is_err());
}
src/config/src/utils/parquet.rs
+15 -5
View File
@@ -22,7 +22,10 @@ use std::{
};
#[cfg(feature = "vortex")]
use arrow::{array::StructArray, datatypes::DataType};
use arrow::{
array::StructArray,
datatypes::{DataType, Field},
};
use arrow::{error::ArrowError, record_batch::RecordBatch};
use arrow_schema::Schema;
#[cfg(feature = "vortex")]
@@ -36,7 +39,7 @@ use parquet::{
#[cfg(feature = "vortex")]
use vortex::{
VortexSessionDefault,
array::{ArrayRef, arrow::IntoArrowArray},
array::{ArrayRef, VortexSessionExecute, arrow::ArrowSessionExt},
buffer::Buffer,
file::OpenOptionsSessionExt,
io::session::RuntimeSessionExt,
@@ -144,11 +147,15 @@ pub type RecordBatchStream = Pin<Box<dyn Stream<Item = Result<RecordBatch, Arrow
/// Convert a single vortex [`ArrayRef`] to an Arrow [`RecordBatch`].
#[cfg(feature = "vortex")]
pub fn vortex_array_to_record_batch(
session: &VortexSession,
array: ArrayRef,
data_type: &DataType,
) -> Result<RecordBatch, ArrowError> {
let array = array
.into_arrow(data_type)
let mut ctx = session.create_execution_ctx();
let target = Field::new("", data_type.clone(), array.dtype().is_nullable());
let array = session
.arrow()
.execute_arrow(array, Some(&target), &mut ctx)
.map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
let struct_array = array
.as_any()
@@ -183,9 +190,12 @@ pub async fn get_recordbatch_reader_from_bytes(
let stream = vortex_stream.then(move |result| {
let arrow_data_type = arrow_data_type.clone();
let session = session.clone();
async move {
match result {
Ok(array) => vortex_array_to_record_batch(array, &arrow_data_type),
Ok(array) => {
vortex_array_to_record_batch(&session, array, &arrow_data_type)
}
Err(e) => Err(ArrowError::ExternalError(Box::new(e))),
}
}
src/handler/grpc/flight/visitor.rs
+9 -35
View File
@@ -53,7 +53,7 @@ impl<'n> TreeNodeVisitor<'n> for RemoteScanVisitor {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
let name = node.name();
if name == "RemoteScanExec" {
let remote_scan_exec = node.as_any().downcast_ref::<RemoteScanExec>().unwrap();
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().unwrap();
self.scan_stats = Some(remote_scan_exec.scan_stats());
Ok(TreeNodeRecursion::Stop)
} else {
@@ -86,7 +86,7 @@ impl<'n> TreeNodeVisitor<'n> for MetricsVisitor {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
let name = node.name();
if name == "RemoteScanExec" {
let remote_scan_exec = node.as_any().downcast_ref::<RemoteScanExec>().unwrap();
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().unwrap();
self.metrics.push(remote_scan_exec.cluster_metrics());
}
Ok(TreeNodeRecursion::Continue)
@@ -114,7 +114,7 @@ impl<'n> TreeNodeVisitor<'n> for PartialErrVisitor {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if node.name() == "RemoteScanExec" {
let remote_scan_exec = node.as_any().downcast_ref::<RemoteScanExec>().unwrap();
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().unwrap();
self.partial_err = Some(remote_scan_exec.partial_err());
Ok(TreeNodeRecursion::Stop)
} else {
@@ -145,7 +145,7 @@ impl<'n> TreeNodeVisitor<'n> for PeakMemoryVisitor {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
let name = node.name();
if name == "RemoteScanExec" {
let remote_scan_exec = node.as_any().downcast_ref::<RemoteScanExec>().unwrap();
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().unwrap();
self.peak_memory = Some(remote_scan_exec.peak_memory());
}
Ok(TreeNodeRecursion::Continue)
@@ -154,40 +154,14 @@ impl<'n> TreeNodeVisitor<'n> for PeakMemoryVisitor {
/// Get the peak memory from the SessionContext's memory pool.
/// This should be called after the query execution to get the peak memory usage.
///
/// # Safety
/// This function uses unsafe code to downcast Arc<dyn MemoryPool> to Arc<PeakMemoryPool>.
/// This is sound because:
/// 1. The memory pool is always created as PeakMemoryPool in datafusion context creation (see
/// exec.rs)
/// 2. Arc::into_raw/from_raw maintain proper Arc semantics
/// 3. The type name is checked at runtime to verify the cast is valid
pub fn get_peak_memory_from_ctx(ctx: &SessionContext) -> Arc<AtomicUsize> {
let memory_pool = ctx.runtime_env().memory_pool.clone();
// Runtime type check: verify this is actually a PeakMemoryPool
let type_name = format!("{:?}", memory_pool);
if !type_name.contains("PeakMemoryPool") {
log::warn!(
"Memory pool is not PeakMemoryPool (type: {}), returning 0",
type_name
);
return Arc::new(AtomicUsize::new(0));
}
// SAFETY: We've verified the concrete type is PeakMemoryPool via runtime check above.
// The memory pool is created as PeakMemoryPool in datafusion context creation.
// We use Arc::into_raw and Arc::from_raw to maintain proper reference counting.
unsafe {
let raw_ptr = Arc::into_raw(memory_pool);
let peak_pool_ptr = raw_ptr as *const PeakMemoryPool;
let peak_memory = (*peak_pool_ptr).peak_memory.clone();
// Reconstruct the Arc to prevent memory leak
let _ = Arc::from_raw(raw_ptr);
peak_memory
}
memory_pool
.downcast_ref::<PeakMemoryPool>()
.unwrap()
.peak_memory
.clone()
}
#[cfg(test)]
src/service/search/datafusion/distributed_plan/codec/aggregate_topk_exec.rs
+3 -3
View File
@@ -40,7 +40,7 @@ pub(crate) fn try_decode(
}
pub(crate) fn try_encode(node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
let Some(node) = node.as_any().downcast_ref::<AggregateTopkExec>() else {
let Some(node) = node.downcast_ref::<AggregateTopkExec>() else {
return internal_err!("Not supported");
};
let plan_node = cluster_rpc::AggregateTopkExecNode {
@@ -120,8 +120,8 @@ mod tests {
let ctx = datafusion::prelude::SessionContext::new();
let plan2 =
physical_plan_from_bytes_with_extension_codec(&plan_bytes, &ctx.task_ctx(), &proto)?;
let plan2 = plan2.as_any().downcast_ref::<AggregateTopkExec>().unwrap();
let plan = plan.as_any().downcast_ref::<AggregateTopkExec>().unwrap();
let plan2 = plan2.downcast_ref::<AggregateTopkExec>().unwrap();
let plan = plan.downcast_ref::<AggregateTopkExec>().unwrap();
// check
assert_eq!(plan.limit(), plan2.limit());
src/service/search/datafusion/distributed_plan/codec/deduplication_exec.rs
+5 -8
View File
@@ -57,14 +57,11 @@ pub(super) fn try_decode(
}
pub(super) fn try_encode(node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
let exec = node
.as_any()
.downcast_ref::<DeduplicationExec>()
.ok_or_else(|| {
datafusion::error::DataFusionError::Internal(
"Failed to downcast to DeduplicationExec".to_string(),
)
})?;
let exec = node.downcast_ref::<DeduplicationExec>().ok_or_else(|| {
datafusion::error::DataFusionError::Internal(
"Failed to downcast to DeduplicationExec".to_string(),
)
})?;
// Encode deduplication columns as PhysicalExprNodes
let mut deduplication_columns_proto = Vec::new();
src/service/search/datafusion/distributed_plan/codec/empty_exec.rs
+3 -3
View File
@@ -62,7 +62,7 @@ pub(crate) fn try_decode(
}
pub(crate) fn try_encode(node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
let Some(node) = node.as_any().downcast_ref::<NewEmptyExec>() else {
let Some(node) = node.downcast_ref::<NewEmptyExec>() else {
return internal_err!("Not supported");
};
let extension_codec = DefaultLogicalExtensionCodec {};
@@ -119,8 +119,8 @@ mod tests {
let ctx = datafusion::prelude::SessionContext::new();
let plan2 =
physical_plan_from_bytes_with_extension_codec(&plan_bytes, &ctx.task_ctx(), &proto)?;
let plan2 = plan2.as_any().downcast_ref::<NewEmptyExec>().unwrap();
let plan = plan.as_any().downcast_ref::<NewEmptyExec>().unwrap();
let plan2 = plan2.downcast_ref::<NewEmptyExec>().unwrap();
let plan = plan.downcast_ref::<NewEmptyExec>().unwrap();
// check
assert_eq!(plan.name(), plan2.name());
src/service/search/datafusion/distributed_plan/codec/enrichment_exec.rs
+3 -3
View File
@@ -42,7 +42,7 @@ pub(crate) fn try_decode(
}
pub(crate) fn try_encode(node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
let Some(node) = node.as_any().downcast_ref::<EnrichmentExec>() else {
let Some(node) = node.downcast_ref::<EnrichmentExec>() else {
return internal_err!("Not supported");
};
let plan_node = cluster_rpc::EnrichmentExecNode {
@@ -92,8 +92,8 @@ mod tests {
let ctx = datafusion::prelude::SessionContext::new();
let plan2 =
physical_plan_from_bytes_with_extension_codec(&plan_bytes, &ctx.task_ctx(), &proto)?;
let plan2 = plan2.as_any().downcast_ref::<EnrichmentExec>().unwrap();
let plan = plan.as_any().downcast_ref::<EnrichmentExec>().unwrap();
let plan2 = plan2.downcast_ref::<EnrichmentExec>().unwrap();
let plan = plan.downcast_ref::<EnrichmentExec>().unwrap();
// check
assert_eq!(plan.trace_id(), plan2.trace_id());
src/service/search/datafusion/distributed_plan/codec/physical_plan_node.rs
+10 -10
View File
@@ -88,25 +88,25 @@ impl PhysicalExtensionCodec for PhysicalPlanNodePhysicalExtensionCodec {
fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
#[cfg(feature = "enterprise")]
if node.as_any().downcast_ref::<NewEmptyExec>().is_some() {
if node.downcast_ref::<NewEmptyExec>().is_some() {
super::empty_exec::try_encode(node, buf)
} else if node.as_any().downcast_ref::<DeduplicationExec>().is_some() {
} else if node.downcast_ref::<DeduplicationExec>().is_some() {
super::deduplication_exec::try_encode(node, buf)
} else if node.as_any().downcast_ref::<AggregateTopkExec>().is_some() {
} else if node.downcast_ref::<AggregateTopkExec>().is_some() {
super::aggregate_topk_exec::try_encode(node, buf)
} else if node.as_any().downcast_ref::<StreamingAggsExec>().is_some() {
} else if node.downcast_ref::<StreamingAggsExec>().is_some() {
super::streaming_aggs_exec::try_encode(node, buf)
} else if node.as_any().downcast_ref::<TmpExec>().is_some() {
} else if node.downcast_ref::<TmpExec>().is_some() {
super::tmp_exec::try_encode(node, buf)
} else if node.as_any().downcast_ref::<EnrichmentExec>().is_some() {
} else if node.downcast_ref::<EnrichmentExec>().is_some() {
super::enrichment_exec::try_encode(node, buf)
} else {
internal_err!("Not supported")
}
#[cfg(not(feature = "enterprise"))]
if node.as_any().downcast_ref::<NewEmptyExec>().is_some() {
if node.downcast_ref::<NewEmptyExec>().is_some() {
super::empty_exec::try_encode(node, buf)
} else if node.as_any().downcast_ref::<DeduplicationExec>().is_some() {
} else if node.downcast_ref::<DeduplicationExec>().is_some() {
super::deduplication_exec::try_encode(node, buf)
} else {
internal_err!("Not supported")
@@ -144,8 +144,8 @@ mod tests {
let ctx = datafusion::prelude::SessionContext::new();
let plan2 =
physical_plan_from_bytes_with_extension_codec(&plan_bytes, &ctx.task_ctx(), &proto)?;
let plan2 = plan2.as_any().downcast_ref::<NewEmptyExec>().unwrap();
let plan = plan.as_any().downcast_ref::<NewEmptyExec>().unwrap();
let plan2 = plan2.downcast_ref::<NewEmptyExec>().unwrap();
let plan = plan.downcast_ref::<NewEmptyExec>().unwrap();
// check
assert_eq!(plan.name(), plan2.name());
src/service/search/datafusion/distributed_plan/codec/streaming_aggs_exec.rs
+8 -11
View File
@@ -36,7 +36,7 @@ pub(crate) fn try_decode(
};
let extension_codec = super::get_physical_extension_codec();
let aggregate_plan = aggregate_plan.try_into_physical_plan(ctx, &extension_codec)?;
let Some(aggregate_plan) = aggregate_plan.as_any().downcast_ref::<AggregateExec>() else {
let Some(aggregate_plan) = aggregate_plan.downcast_ref::<AggregateExec>() else {
return internal_err!("aggregate_plan is not an AggregateExec");
};
let aggregate_plan = Arc::new(aggregate_plan.clone());
@@ -61,7 +61,7 @@ pub(crate) fn try_decode(
}
pub(crate) fn try_encode(node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
let Some(node) = node.as_any().downcast_ref::<StreamingAggsExec>() else {
let Some(node) = node.downcast_ref::<StreamingAggsExec>() else {
return internal_err!("Not supported");
};
@@ -169,8 +169,8 @@ mod tests {
let ctx = datafusion::prelude::SessionContext::new();
let plan2 =
physical_plan_from_bytes_with_extension_codec(&plan_bytes, &ctx.task_ctx(), &proto)?;
let plan2 = plan2.as_any().downcast_ref::<StreamingAggsExec>().unwrap();
let plan = plan.as_any().downcast_ref::<StreamingAggsExec>().unwrap();
let plan2 = plan2.downcast_ref::<StreamingAggsExec>().unwrap();
let plan = plan.downcast_ref::<StreamingAggsExec>().unwrap();
// check
assert_eq!(plan.id(), plan2.id());
@@ -255,8 +255,8 @@ mod tests {
ctx.register_udf(STR_MATCH_UDF.clone());
let plan2 =
physical_plan_from_bytes_with_extension_codec(&plan_bytes, &ctx.task_ctx(), &proto)?;
let plan2 = plan2.as_any().downcast_ref::<StreamingAggsExec>().unwrap();
let plan = plan.as_any().downcast_ref::<StreamingAggsExec>().unwrap();
let plan2 = plan2.downcast_ref::<StreamingAggsExec>().unwrap();
let plan = plan.downcast_ref::<StreamingAggsExec>().unwrap();
// check
assert_eq!(plan.id(), plan2.id());
@@ -268,14 +268,11 @@ mod tests {
// Verify that the aggregate plan contains a FilterExec
let agg_plan_decoded = plan2.aggregate_plan();
let agg_exec = agg_plan_decoded
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let agg_exec = agg_plan_decoded.as_ref();
let input_plan = agg_exec.input();
// The input should be a FilterExec
assert!(input_plan.as_any().downcast_ref::<FilterExec>().is_some());
assert!(input_plan.downcast_ref::<FilterExec>().is_some());
Ok(())
}
src/service/search/datafusion/distributed_plan/codec/tmp_exec.rs
+3 -3
View File
@@ -42,7 +42,7 @@ pub(crate) fn try_decode(
}
pub(crate) fn try_encode(node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
let Some(node) = node.as_any().downcast_ref::<TmpExec>() else {
let Some(node) = node.downcast_ref::<TmpExec>() else {
return internal_err!("Not supported");
};
let plan_node = cluster_rpc::TmpExecNode {
@@ -111,8 +111,8 @@ mod tests {
let ctx = datafusion::prelude::SessionContext::new();
let plan2 =
physical_plan_from_bytes_with_extension_codec(&plan_bytes, &ctx.task_ctx(), &proto)?;
let plan2 = plan2.as_any().downcast_ref::<TmpExec>().unwrap();
let plan = plan.as_any().downcast_ref::<TmpExec>().unwrap();
let plan2 = plan2.downcast_ref::<TmpExec>().unwrap();
let plan = plan.downcast_ref::<TmpExec>().unwrap();
// check
assert_eq!(plan.trace_id(), plan2.trace_id());
src/service/search/datafusion/distributed_plan/distribute_analyze_exec.rs
+1 -6
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use arrow::{
array::{Int32Builder, StringBuilder},
@@ -104,11 +104,6 @@ impl ExecutionPlan for DistributeAnalyzeExec {
"DistributeAnalyzeExec"
}
/// Return a reference to Any that can be used for downcasting
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.cache
}
src/service/search/datafusion/distributed_plan/empty_exec.rs
+3 -10
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use arrow_schema::SortOptions;
use config::TIMESTAMP_COL_NAME;
@@ -196,11 +196,6 @@ impl ExecutionPlan for NewEmptyExec {
"NewEmptyExec"
}
/// Return a reference to Any that can be used for downcasting
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.cache
}
@@ -314,12 +309,10 @@ mod tests {
#[test]
fn test_as_any_downcasts() {
use datafusion::physical_plan::ExecutionPlan;
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
let exec = NewEmptyExec::new("t", schema.clone(), None, &[], None, false, schema);
let any = exec.as_any();
assert!(any.downcast_ref::<NewEmptyExec>().is_some());
let exec_plan: &dyn ExecutionPlan = &exec;
assert!(exec_plan.downcast_ref::<NewEmptyExec>().is_some());
}
#[test]
src/service/search/datafusion/distributed_plan/enrich_exec.rs
+3 -10
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use config::utils::record_batch_ext::convert_json_to_record_batch;
use datafusion::{
@@ -93,11 +93,6 @@ impl ExecutionPlan for EnrichExec {
"EnrichExec"
}
/// Return a reference to Any that can be used for downcasting
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.cache
}
@@ -222,12 +217,10 @@ mod tests {
#[test]
fn test_as_any_downcasts() {
use datafusion::physical_plan::ExecutionPlan;
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
let exec = EnrichExec::new("org", "stream", schema);
let any = exec.as_any();
assert!(any.downcast_ref::<EnrichExec>().is_some());
let exec_plan: &dyn ExecutionPlan = &exec;
assert!(exec_plan.downcast_ref::<EnrichExec>().is_some());
}
#[test]
src/service/search/datafusion/distributed_plan/enrichment_exec.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use arrow::{
array::{ArrayRef, RecordBatch},
@@ -146,10 +146,6 @@ impl ExecutionPlan for EnrichmentExec {
"EnrichmentExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.cache
}
src/service/search/datafusion/distributed_plan/mod.rs
+1 -5
View File
@@ -54,7 +54,6 @@ impl NewEmptyExecVisitor {
self.plan
.as_ref()
.unwrap()
.as_any()
.downcast_ref::<NewEmptyExec>()
.unwrap()
}
@@ -107,7 +106,7 @@ impl TreeNodeRewriter for ReplaceTableScanExec {
#[cfg(test)]
mod tests {
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use datafusion::physical_plan::{ExecutionPlan, PlanProperties};
@@ -128,9 +127,6 @@ mod tests {
fn name(&self) -> &'static str {
"NewEmptyExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
panic!("not needed")
}
src/service/search/datafusion/distributed_plan/remote_scan_exec.rs
+3 -10
View File
@@ -13,10 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{
any::Any,
sync::{Arc, atomic::AtomicUsize},
};
use std::sync::{Arc, atomic::AtomicUsize};
use arrow_schema::SchemaRef;
use config::{
@@ -207,11 +204,6 @@ impl ExecutionPlan for RemoteScanExec {
"RemoteScanExec"
}
/// Return a reference to Any that can be used for downcasting
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
Arc::clone(&self.input.schema())
}
@@ -493,7 +485,8 @@ mod tests {
#[test]
fn test_remote_scan_exec_as_any() {
let exec = make_remote_exec();
assert!(exec.as_any().downcast_ref::<RemoteScanExec>().is_some());
let exec_plan: &dyn ExecutionPlan = &exec;
assert!(exec_plan.downcast_ref::<RemoteScanExec>().is_some());
}
#[test]
src/service/search/datafusion/distributed_plan/rewrite.rs
+1 -1
View File
@@ -75,7 +75,7 @@ impl TreeNodeRewriter for TantivyOptimizeRewriter {
fn f_up(&mut self, node: Arc<dyn ExecutionPlan>) -> Result<Transformed<Self::Node>> {
if node.name() == "AggregateExec" {
let aggregate = node.as_any().downcast_ref::<AggregateExec>().unwrap();
let aggregate = node.downcast_ref::<AggregateExec>().unwrap();
if *aggregate.mode() == AggregateMode::Partial {
let tantivy_exec = Arc::new(TantivyOptimizeExec::new(
self.query.clone(),
src/service/search/datafusion/exec.rs
+38 -4
View File
@@ -60,7 +60,10 @@ use super::{
storage::file_list, udf::transform_udf::get_all_transform,
};
use crate::service::search::{
datafusion::table_provider::{listing_adapter::ListingTableAdapter, uniontable::NewUnionTable},
datafusion::{
storage::file_statistics_cache,
table_provider::{listing_adapter::ListingTableAdapter, uniontable::NewUnionTable},
},
index::IndexCondition,
};
@@ -107,6 +110,29 @@ pub fn create_session_config(
.execution
.skip_physical_aggregate_schema_check = true;
// DataFusion 54 executes uncorrelated scalar subqueries physically via
// `ScalarSubqueryExec`/`ScalarSubqueryExpr` instead of rewriting them into joins.
// `ScalarSubqueryExpr` can only be (de)serialized inside its surrounding
// `ScalarSubqueryExec`, which breaks our distributed plan splitting across the Flight
// boundary. Disable the physical path so `ScalarSubqueryToJoin` decorrelates them into
// joins again, keeping the serialized follower plans valid.
config
.options_mut()
.optimizer
.enable_physical_uncorrelated_scalar_subquery = false;
// DataFusion 54 builds a runtime `DynamicFilterPhysicalExpr` from a `HashJoinExec`'s
// build-side join keys and pushes it into the probe-side scan. That runtime state can't
// cross our distributed RemoteScan/Flight boundary, and after our custom join rewrites
// (`swap_inputs` + broadcast/enrichment join) the filter ends up referencing build-side
// columns by index against the projected probe-side batch, producing
// "PhysicalExpr Column references column ... but input schema only has N columns" at
// execution time. Disable join dynamic filter pushdown to keep the split plans valid.
config
.options_mut()
.optimizer
.enable_join_dynamic_filter_pushdown = false;
Ok(config)
}
@@ -126,9 +152,9 @@ pub async fn create_runtime_env(trace_id: &str, memory_limit: usize) -> Result<R
RuntimeEnvBuilder::new().with_object_store_registry(Arc::new(object_store_registry));
if cfg.limit.datafusion_file_stat_cache_max_size > 0 {
let cache_config = CacheManagerConfig::default();
let cache_config = cache_config.with_files_statistics_cache(Some(
super::storage::file_statistics_cache::GLOBAL_CACHE.clone(),
));
let cache_config = cache_config
.with_file_statistics_cache(Some(file_statistics_cache::GLOBAL_CACHE.clone()))
.with_file_statistics_cache_limit(cfg.limit.datafusion_file_stat_cache_max_size);
builder = builder.with_cache_manager(cache_config);
}
@@ -614,6 +640,14 @@ mod tests {
assert_eq!(config.options().sql_parser.dialect, Dialect::PostgreSQL);
assert!(!config.options().execution.listing_table_ignore_subdirectory);
assert!(config.information_schema());
// Join dynamic filter pushdown must stay disabled: its runtime filter can't cross our
// distributed RemoteScan/Flight boundary and breaks our custom join rewrites.
assert!(
!config
.options()
.optimizer
.enable_join_dynamic_filter_pushdown
);
Ok(())
}
src/service/search/datafusion/optimizer/logical_optimizer/limit_join_right_side.rs
+8 -9
View File
@@ -375,15 +375,14 @@ mod tests {
let expected = vec![
"HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(prod_id@1, prod_id@0)], projection=[usr_id@0, prod_id@1]",
" ProjectionExec: expr=[usr_id@1 as usr_id, prod_id@0 as prod_id]",
" HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(usr_id@0, usr_id@0)], projection=[prod_id@1, usr_id@2]",
" ProjectionExec: expr=[usr_id@0 as usr_id, prod_id@1 as prod_id]",
" DeduplicationExec: columns: [usr_id@0]",
" SortExec: expr=[usr_id@0 DESC NULLS LAST, _timestamp@2 DESC NULLS LAST], preserve_partitioning=[false]",
" SortPreservingMergeExec: [_timestamp@2 DESC NULLS LAST], fetch=50000",
" SortExec: TopK(fetch=50000), expr=[_timestamp@2 DESC NULLS LAST], preserve_partitioning=[true]",
" DataSourceExec: partitions=2, partition_sizes=[1, 1]",
" DataSourceExec: partitions=1, partition_sizes=[1]",
" HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(usr_id@0, usr_id@0)], projection=[usr_id@2, prod_id@1]",
" ProjectionExec: expr=[usr_id@0 as usr_id, prod_id@1 as prod_id]",
" DeduplicationExec: columns: [usr_id@0]",
" SortExec: expr=[usr_id@0 DESC NULLS LAST, _timestamp@2 DESC NULLS LAST], preserve_partitioning=[false]",
" SortPreservingMergeExec: [_timestamp@2 DESC NULLS LAST], fetch=50000",
" SortExec: TopK(fetch=50000), expr=[_timestamp@2 DESC NULLS LAST], preserve_partitioning=[true]",
" DataSourceExec: partitions=2, partition_sizes=[1, 1]",
" DataSourceExec: partitions=1, partition_sizes=[1]",
" ProjectionExec: expr=[prod_id@0 as prod_id]",
" DeduplicationExec: columns: [prod_id@0]",
" SortExec: expr=[prod_id@0 DESC NULLS LAST, _timestamp@1 DESC NULLS LAST], preserve_partitioning=[false]",
src/service/search/datafusion/optimizer/physical_optimizer/distribute_analyze.rs
+4 -3
View File
@@ -29,7 +29,7 @@ use crate::service::search::datafusion::distributed_plan::{
// replace the AnalyzeExec to DistributeAnalyzeExec
pub fn optimize_distribute_analyze(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
if let Some(analyze) = plan.as_any().downcast_ref::<AnalyzeExec>() {
if let Some(analyze) = plan.downcast_ref::<AnalyzeExec>() {
let distribute_analyze = Arc::new(DistributeAnalyzeExec::new(
analyze.verbose(),
analyze.show_statistics(),
@@ -55,7 +55,7 @@ impl TreeNodeRewriter for AnalyzeRewrite {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: Arc<dyn ExecutionPlan>) -> Result<Transformed<Self::Node>> {
if let Some(remote) = node.as_any().downcast_ref::<RemoteScanExec>() {
if let Some(remote) = node.downcast_ref::<RemoteScanExec>() {
let remote = Arc::new(remote.clone().set_analyze()) as Arc<dyn ExecutionPlan>;
return Ok(Transformed::yes(remote));
}
@@ -84,6 +84,7 @@ mod tests {
false,
false,
vec![],
None,
Arc::new(EmptyExec::new(get_schema())),
get_schema(),
));
@@ -102,7 +103,7 @@ mod tests {
) as Arc<dyn ExecutionPlan>;
let mut rewriter = AnalyzeRewrite::new();
let plan = plan.rewrite(&mut rewriter).unwrap().data;
let remote_scan = plan.as_any().downcast_ref::<RemoteScanExec>().unwrap();
let remote_scan = plan.downcast_ref::<RemoteScanExec>().unwrap();
assert!(remote_scan.analyze());
}
}
src/service/search/datafusion/optimizer/physical_optimizer/enrichment.rs
+4 -4
View File
@@ -88,7 +88,7 @@ impl TreeNodeRewriter for EnrichmentExecRewriter {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: Arc<dyn ExecutionPlan>) -> Result<Transformed<Self::Node>> {
if let Some(new_empty) = node.as_any().downcast_ref::<NewEmptyExec>() {
if let Some(new_empty) = node.downcast_ref::<NewEmptyExec>() {
let table_name = TableReference::parse_str(new_empty.name());
if matches!(
table_name.schema(),
@@ -140,7 +140,7 @@ pub fn should_use_enrichment_broadcast_join(plan: &Arc<dyn ExecutionPlan>) -> bo
let mut result = false;
let _ = plan.apply(|node| {
Ok(if node.name() == "HashJoinExec" {
if let Some(hash_join) = node.as_any().downcast_ref::<HashJoinExec>() {
if let Some(hash_join) = node.downcast_ref::<HashJoinExec>() {
result = is_enrichment_table(hash_join.left())
&& is_valid_enrich_broadcast_right(hash_join.right());
}
@@ -157,7 +157,7 @@ pub fn is_enrichment_table(plan: &Arc<dyn ExecutionPlan>) -> bool {
let mut found_enrichment = false;
let _ = plan.apply(|node| {
Ok(
if let Some(new_empty) = node.as_any().downcast_ref::<NewEmptyExec>() {
if let Some(new_empty) = node.downcast_ref::<NewEmptyExec>() {
let table_name = TableReference::parse_str(new_empty.name());
if matches!(
table_name.schema(),
@@ -180,7 +180,7 @@ fn is_valid_enrich_broadcast_right(plan: &Arc<dyn ExecutionPlan>) -> bool {
let mut is_valid = true;
let _ = plan.apply(|node| {
Ok(
if let Some(new_empty) = node.as_any().downcast_ref::<NewEmptyExec>() {
if let Some(new_empty) = node.downcast_ref::<NewEmptyExec>() {
let table_name = TableReference::parse_str(new_empty.name());
if matches!(
table_name.schema(),
src/service/search/datafusion/optimizer/physical_optimizer/index.rs
+5 -5
View File
@@ -177,7 +177,7 @@ impl TreeNodeRewriter for IndexOptimizer {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
if let Some(filter) = node.as_any().downcast_ref::<FilterExec>() {
if let Some(filter) = node.downcast_ref::<FilterExec>() {
self.has_filter = true;
let mut index_conditions = IndexCondition::new();
let mut other_conditions = Vec::new();
@@ -254,7 +254,7 @@ fn construct_filter_exec(
// Check if the expression is valid for the index.
fn is_expr_valid_for_index(expr: &Arc<dyn PhysicalExpr>, index_fields: &HashSet<String>) -> bool {
if let Some(expr) = expr.as_any().downcast_ref::<BinaryExpr>() {
if let Some(expr) = expr.downcast_ref::<BinaryExpr>() {
match expr.op() {
Operator::Eq | Operator::NotEq => {
let column = if is_value(expr.left()) && is_column(expr.right()) {
@@ -275,7 +275,7 @@ fn is_expr_valid_for_index(expr: &Arc<dyn PhysicalExpr>, index_fields: &HashSet<
}
_ => return false,
}
} else if let Some(expr) = expr.as_any().downcast_ref::<InListExpr>() {
} else if let Some(expr) = expr.downcast_ref::<InListExpr>() {
if !is_column(expr.expr()) || !index_fields.contains(get_column_name(expr.expr())) {
return false;
}
@@ -285,7 +285,7 @@ fn is_expr_valid_for_index(expr: &Arc<dyn PhysicalExpr>, index_fields: &HashSet<
return false;
}
}
} else if let Some(expr) = expr.as_any().downcast_ref::<ScalarFunctionExpr>() {
} else if let Some(expr) = expr.downcast_ref::<ScalarFunctionExpr>() {
let name = expr.name();
return match name {
MATCH_ALL_UDF_NAME => {
@@ -303,7 +303,7 @@ fn is_expr_valid_for_index(expr: &Arc<dyn PhysicalExpr>, index_fields: &HashSet<
}
_ => false,
};
} else if let Some(expr) = expr.as_any().downcast_ref::<NotExpr>() {
} else if let Some(expr) = expr.downcast_ref::<NotExpr>() {
return is_expr_valid_for_index(expr.arg(), index_fields);
} else {
return false;
src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/count.rs
+1 -1
View File
@@ -69,7 +69,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleCountVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if let Some(aggregate) = node.as_any().downcast_ref::<AggregateExec>() {
if let Some(aggregate) = node.downcast_ref::<AggregateExec>() {
if aggregate.group_expr().is_empty()
&& aggregate.aggr_expr().len() == 1
&& is_count_rows_aggregate(&aggregate.aggr_expr()[0])
src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/distinct.rs
+5 -5
View File
@@ -82,7 +82,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleDistinctVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if let Some(sort_merge) = node.as_any().downcast_ref::<SortPreservingMergeExec>() {
if let Some(sort_merge) = node.downcast_ref::<SortPreservingMergeExec>() {
if let Some(fetch) = sort_merge.fetch()
&& fetch > 0
&& sort_merge.expr().len() == 1
@@ -99,7 +99,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleDistinctVisitor {
// if not simple distinct, stop visiting
self.simple_distinct = None;
return Ok(TreeNodeRecursion::Stop);
} else if let Some(aggregate) = node.as_any().downcast_ref::<AggregateExec>() {
} else if let Some(aggregate) = node.downcast_ref::<AggregateExec>() {
// check if the group by is simple distinct,
// only one group by field, no aggregate function
if aggregate.group_expr().expr().len() == 1
@@ -122,7 +122,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleDistinctVisitor {
// if not simple distinct, stop visiting
self.simple_distinct = None;
return Ok(TreeNodeRecursion::Stop);
} else if let Some(projection) = node.as_any().downcast_ref::<ProjectionExec>() {
} else if let Some(projection) = node.downcast_ref::<ProjectionExec>() {
// Check ProjectionExec for the structure: [index_field]
let exprs = projection.expr();
if exprs.len() == 1 {
@@ -133,7 +133,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleDistinctVisitor {
// If projection doesn't have exactly 2 expressions, stop visiting
self.simple_distinct = None;
return Ok(TreeNodeRecursion::Stop);
} else if let Some(filter) = node.as_any().downcast_ref::<FilterExec>() {
} else if let Some(filter) = node.downcast_ref::<FilterExec>() {
let predicate = filter.predicate();
let exprs = split_conjunction(predicate);
if exprs.len() == 2 && is_only_timestamp_filter(&exprs) {
@@ -159,7 +159,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleDistinctVisitor {
}
fn is_simple_str_match(expr: &Arc<dyn PhysicalExpr>) -> Option<String> {
if let Some(func) = expr.as_any().downcast_ref::<ScalarFunctionExpr>()
if let Some(func) = expr.downcast_ref::<ScalarFunctionExpr>()
&& func.fun().name().to_lowercase() == "str_match"
&& func.args().len() == 2
{
src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/histogram.rs
+9 -9
View File
@@ -86,7 +86,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleHistogramVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if let Some(aggregate) = node.as_any().downcast_ref::<AggregateExec>() {
if let Some(aggregate) = node.downcast_ref::<AggregateExec>() {
// Check if the AggregateExec matches SimpleHistogram pattern
if aggregate.group_expr().expr().len() == 1
&& aggregate.aggr_expr().len() == 1
@@ -118,7 +118,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleHistogramVisitor {
// If AggregateExec doesn't match SimpleHistogram pattern, stop visiting
self.simple_histogram = None;
return Ok(TreeNodeRecursion::Stop);
} else if let Some(projection) = node.as_any().downcast_ref::<ProjectionExec>() {
} else if let Some(projection) = node.downcast_ref::<ProjectionExec>() {
// Check ProjectionExec for the structure: [histogram(_timestamp), count(*)]
let exprs = projection.expr();
if exprs.len() == 2 {
@@ -139,7 +139,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleHistogramVisitor {
}
fn get_data_bin(expr: &Arc<dyn PhysicalExpr>) -> Option<&ScalarFunctionExpr> {
if let Some(func) = expr.as_any().downcast_ref::<ScalarFunctionExpr>()
if let Some(func) = expr.downcast_ref::<ScalarFunctionExpr>()
&& func.fun().name().to_lowercase() == "date_bin"
{
Some(func)
@@ -150,7 +150,7 @@ fn get_data_bin(expr: &Arc<dyn PhysicalExpr>) -> Option<&ScalarFunctionExpr> {
// unit: microseconds
fn get_histogram_interval(expr: &Arc<dyn PhysicalExpr>) -> Option<u64> {
let interval = expr.as_any().downcast_ref::<Literal>()?.value();
let interval = expr.downcast_ref::<Literal>()?.value();
match interval {
ScalarValue::IntervalMonthDayNano(Some(interval)) => {
// convert interval to nanoseconds
@@ -168,16 +168,16 @@ fn get_histogram_interval(expr: &Arc<dyn PhysicalExpr>) -> Option<u64> {
/// `to_timestamp_micros(_timestamp + offset)` — the shape histogram() with a timezone
/// rewrites to — yields the offset. None when the source is not the timestamp column.
fn get_timestamp_offset(expr: &Arc<dyn PhysicalExpr>) -> Option<i64> {
let func = expr.as_any().downcast_ref::<ScalarFunctionExpr>()?;
let func = expr.downcast_ref::<ScalarFunctionExpr>()?;
let arg = func.args().first()?;
if get_column_name(arg) == TIMESTAMP_COL_NAME {
return Some(0);
}
let bin = arg.as_any().downcast_ref::<BinaryExpr>()?;
let bin = arg.downcast_ref::<BinaryExpr>()?;
if *bin.op() != Operator::Plus || get_column_name(bin.left()) != TIMESTAMP_COL_NAME {
return None;
}
match bin.right().as_any().downcast_ref::<Literal>()?.value() {
match bin.right().downcast_ref::<Literal>()?.value() {
ScalarValue::Int64(Some(ts_offset)) => Some(*ts_offset),
_ => None,
}
@@ -239,7 +239,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleMultiHistogramVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if let Some(aggregate) = node.as_any().downcast_ref::<AggregateExec>() {
if let Some(aggregate) = node.downcast_ref::<AggregateExec>() {
// Exactly 2 group-by expressions (histogram + breakdown) and 1 aggregate (count(*))
if aggregate.group_expr().expr().len() == 2
&& aggregate.aggr_expr().len() == 1
@@ -285,7 +285,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleMultiHistogramVisitor {
// If AggregateExec doesn't match, stop visiting
self.simple_multi_histogram = None;
return Ok(TreeNodeRecursion::Stop);
} else if let Some(projection) = node.as_any().downcast_ref::<ProjectionExec>() {
} else if let Some(projection) = node.downcast_ref::<ProjectionExec>() {
// Projection should have 3 expressions: timestamp, breakdown, count
let exprs = projection.expr();
if exprs.len() == 3 {
src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs
+6 -18
View File
@@ -132,11 +132,7 @@ impl TreeNodeRewriter for FollowerIndexOptimizer {
return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop));
}
if plan
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.is_some()
{
if plan.downcast_ref::<SortPreservingMergeExec>().is_some() {
// Check for SimpleSelect
if let Some(index_optimize_mode) = is_simple_select(Arc::clone(&plan)) {
*self.index_optimizer_mode.lock() = Some(index_optimize_mode);
@@ -158,7 +154,7 @@ impl TreeNodeRewriter for FollowerIndexOptimizer {
}
}
return Ok(Transformed::new(plan, false, TreeNodeRecursion::Continue));
} else if plan.as_any().downcast_ref::<AggregateExec>().is_some() {
} else if plan.downcast_ref::<AggregateExec>().is_some() {
// Check for SimpleCount
if let Some(index_optimize_mode) = is_simple_count(Arc::clone(&plan)) {
*self.index_optimizer_mode.lock() = Some(index_optimize_mode);
@@ -242,11 +238,7 @@ impl TreeNodeRewriter for LeaderIndexOptimizer {
return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop));
}
if plan
.as_any()
.downcast_ref::<SortPreservingMergeExec>()
.is_some()
{
if plan.downcast_ref::<SortPreservingMergeExec>().is_some() {
// Get the index fields of the underlying table
let mut visitor = TableNameVisitor::new();
plan.visit(&mut visitor)?;
@@ -298,7 +290,7 @@ impl TreeNodeRewriter for IndexOptimizerRewrite {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: Arc<dyn ExecutionPlan>) -> Result<Transformed<Self::Node>> {
if let Some(remote) = node.as_any().downcast_ref::<RemoteScanExec>() {
if let Some(remote) = node.downcast_ref::<RemoteScanExec>() {
let remote = Arc::new(
remote
.clone()
@@ -327,7 +319,7 @@ impl<'n> TreeNodeVisitor<'n> for TableNameVisitor {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
let name = node.name();
if name == "NewEmptyExec" {
let table = node.as_any().downcast_ref::<NewEmptyExec>().unwrap();
let table = node.downcast_ref::<NewEmptyExec>().unwrap();
self.table_name = Some(TableReference::from(table.name()));
Ok(TreeNodeRecursion::Stop)
} else {
@@ -402,11 +394,7 @@ mod tests {
let result = plan.rewrite(&mut rewriter).unwrap();
assert!(result.transformed, "plan should be marked as transformed");
assert_eq!(result.data.name(), "RemoteScanExec");
let remote_scan = result
.data
.as_any()
.downcast_ref::<RemoteScanExec>()
.unwrap();
let remote_scan = result.data.downcast_ref::<RemoteScanExec>().unwrap();
assert_eq!(remote_scan.index_optimize_mode(), Some(mode));
}
src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/select.rs
+1 -1
View File
@@ -64,7 +64,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleSelectVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if let Some(sort_merge) = node.as_any().downcast_ref::<SortPreservingMergeExec>() {
if let Some(sort_merge) = node.downcast_ref::<SortPreservingMergeExec>() {
if let Some(fetch) = sort_merge.fetch()
&& fetch > 0
&& sort_merge.expr().len() == 1
src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/topn.rs
+3 -3
View File
@@ -89,7 +89,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleTopnVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if let Some(sort_merge) = node.as_any().downcast_ref::<SortPreservingMergeExec>() {
if let Some(sort_merge) = node.downcast_ref::<SortPreservingMergeExec>() {
// need a fetch limit, a primary sort plus up to one secondary sort per group field
let fetch = sort_merge.fetch().unwrap_or(0);
let sort_exprs = sort_merge.expr();
@@ -120,7 +120,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleTopnVisitor {
// fields are filled in at the AggregateExec
self.simple_topn = Some((vec![], fetch, !first_sort.options.descending));
} else if let Some(projection) = node.as_any().downcast_ref::<ProjectionExec>() {
} else if let Some(projection) = node.downcast_ref::<ProjectionExec>() {
// expect [field1, .., fieldN, count(*)]; the exact length is validated against
// the group by at the AggregateExec
let exprs = projection.expr();
@@ -141,7 +141,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleTopnVisitor {
{
return self.reject();
}
} else if let Some(aggregate) = node.as_any().downcast_ref::<AggregateExec>() {
} else if let Some(aggregate) = node.downcast_ref::<AggregateExec>() {
let group_len = aggregate.group_expr().expr().len();
if !(1..=MAX_SIMPLE_TOPN_FIELDS).contains(&group_len)
|| aggregate.aggr_expr().len() != 1
src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/utils.rs
+3 -3
View File
@@ -83,7 +83,7 @@ pub(crate) mod tests {
let mut visitor = AggregateVisitor::new();
let _ = plan.visit(&mut visitor);
let data = visitor.get_data();
data.map(|v| v.as_any().downcast_ref::<AggregateExec>().unwrap().clone())
data.map(|v| v.downcast_ref::<AggregateExec>().unwrap().clone())
}
struct AggregateVisitor {
@@ -111,7 +111,7 @@ pub(crate) mod tests {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if node.name() == "AggregateExec" {
let agg = node.as_any().downcast_ref::<AggregateExec>().unwrap();
let agg = node.downcast_ref::<AggregateExec>().unwrap();
if *agg.mode() == AggregateMode::Partial {
self.data = Some(node.clone());
Ok(TreeNodeRecursion::Stop)
@@ -155,7 +155,7 @@ pub(crate) mod tests {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if node.name() == "RemoteScanExec" {
let remote_scan = node.as_any().downcast_ref::<RemoteScanExec>().unwrap();
let remote_scan = node.downcast_ref::<RemoteScanExec>().unwrap();
self.data.push(Arc::new(remote_scan.clone()));
}
Ok(TreeNodeRecursion::Continue)
src/service/search/datafusion/optimizer/physical_optimizer/join_reorder.rs
+1 -1
View File
@@ -62,7 +62,7 @@ impl PhysicalOptimizerRule for JoinReorderRule {
}
fn swap_join_order(plan: Arc<dyn ExecutionPlan>) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
if let Some(hash_join) = plan.as_any().downcast_ref::<HashJoinExec>() {
if let Some(hash_join) = plan.downcast_ref::<HashJoinExec>() {
let left = hash_join.left();
let right = hash_join.right();
src/service/search/datafusion/optimizer/physical_optimizer/remote_scan.rs
+9 -10
View File
@@ -273,7 +273,7 @@ impl TreeNodeRewriter for RemoteScanRewriter {
// For sort, we should add a SortPreservingMergeExec
if child.name() == "SortExec" {
let table_name = visitor.table_name.clone().unwrap();
let sort = child.as_any().downcast_ref::<SortExec>().unwrap();
let sort = child.downcast_ref::<SortExec>().unwrap();
let sort_merge = Arc::new(
SortPreservingMergeExec::new(
LexOrdering::new(sort.expr().to_vec()).unwrap(),
@@ -338,7 +338,7 @@ fn wrap_partial_reduce(
if !partial_reduce_enabled {
return Ok(input);
}
let Some(agg) = input.as_any().downcast_ref::<AggregateExec>() else {
let Some(agg) = input.downcast_ref::<AggregateExec>() else {
return Ok(input);
};
if *agg.mode() != AggregateMode::Partial {
@@ -429,7 +429,7 @@ impl<'n> TreeNodeVisitor<'n> for TableNameVisitor {
self.has_remote_scan = true;
Ok(TreeNodeRecursion::Stop)
} else if name == "NewEmptyExec" {
let table = node.as_any().downcast_ref::<NewEmptyExec>().unwrap();
let table = node.downcast_ref::<NewEmptyExec>().unwrap();
self.table_name = Some(TableReference::from(table.name()));
Ok(TreeNodeRecursion::Continue)
} else {
@@ -546,7 +546,7 @@ mod tests {
)?;
let input: Arc<dyn ExecutionPlan> = Arc::new(agg);
let result = wrap_partial_reduce(true, input)?;
let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
let result_agg = result.downcast_ref::<AggregateExec>().unwrap();
assert_eq!(*result_agg.mode(), AggregateMode::Single);
Ok(())
}
@@ -566,7 +566,7 @@ mod tests {
)?;
let input: Arc<dyn ExecutionPlan> = Arc::new(agg);
let result = wrap_partial_reduce(true, input)?;
let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
let result_agg = result.downcast_ref::<AggregateExec>().unwrap();
assert_eq!(*result_agg.mode(), AggregateMode::PartialReduce);
assert!(result_agg.group_expr().expr().is_empty());
Ok(())
@@ -592,13 +592,12 @@ mod tests {
)?;
let input: Arc<dyn ExecutionPlan> = Arc::new(agg);
let result = wrap_partial_reduce(true, input)?;
let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
let result_agg = result.downcast_ref::<AggregateExec>().unwrap();
assert_eq!(*result_agg.mode(), AggregateMode::PartialReduce);
assert_eq!(result_agg.group_expr().expr().len(), 1);
let (expr, name) = &result_agg.group_expr().expr()[0];
assert_eq!(name, "a");
let col_expr = expr
.as_any()
.downcast_ref::<PhysicalColumn>()
.expect("group expr should be a Column reference into partial output schema");
assert_eq!(col_expr.name(), "a");
@@ -638,7 +637,7 @@ mod tests {
)?;
let input: Arc<dyn ExecutionPlan> = Arc::new(agg);
let result = wrap_partial_reduce(true, input)?;
let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
let result_agg = result.downcast_ref::<AggregateExec>().unwrap();
assert_eq!(*result_agg.mode(), AggregateMode::PartialReduce);
// The critical invariant: input_schema must be the original scan schema so that
// aggregate expression argument types (e.g. _timestamp:Int64) are reconstructed
@@ -694,7 +693,7 @@ mod tests {
)?;
let input: Arc<dyn ExecutionPlan> = Arc::new(agg);
let result = wrap_partial_reduce(true, input)?;
let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
let result_agg = result.downcast_ref::<AggregateExec>().unwrap();
assert_eq!(*result_agg.mode(), AggregateMode::PartialReduce);
let null_exprs = result_agg.group_expr().null_expr();
@@ -702,7 +701,7 @@ mod tests {
// Both must still be Literal(NULL), NOT remapped to Column refs.
for (expr, _) in null_exprs {
assert!(
expr.as_any().is::<Literal>(),
expr.is::<Literal>(),
"null_expr must remain a Literal, not be replaced with a Column ref"
);
}
src/service/search/datafusion/optimizer/physical_optimizer/rewrite_match.rs
+9 -12
View File
@@ -99,7 +99,7 @@ impl TreeNodeRewriter for PlanRewriter {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, plan: Self::Node) -> Result<Transformed<Self::Node>> {
if let Some(filter) = plan.as_any().downcast_ref::<FilterExec>() {
if let Some(filter) = plan.downcast_ref::<FilterExec>() {
let predicate = filter.predicate();
// 1. Check if the predicate contains any match_all functions
@@ -160,7 +160,7 @@ impl TreeNodeRewriter for ColumnIndexRewriter {
fn f_up(&mut self, expr: Self::Node) -> Result<Transformed<Self::Node>> {
// Check if this is a Column expression
if let Some(column) = expr.as_any().downcast_ref::<Column>() {
if let Some(column) = expr.downcast_ref::<Column>() {
let field_name = column.name();
// Find the new index in the new schema
if let Ok(new_index) = self.new_schema.index_of(field_name) {
@@ -210,12 +210,9 @@ fn rewrite_match_all_physical(
schema: SchemaRef,
fields: &[(String, DataType)],
) -> Result<Arc<dyn PhysicalExpr>> {
let scalar_fn = expr
.as_any()
.downcast_ref::<ScalarFunctionExpr>()
.ok_or_else(|| {
DataFusionError::Internal("Expected ScalarFunctionExpr for match_all".to_string())
})?;
let scalar_fn = expr.downcast_ref::<ScalarFunctionExpr>().ok_or_else(|| {
DataFusionError::Internal("Expected ScalarFunctionExpr for match_all".to_string())
})?;
let name = scalar_fn.name();
let args = scalar_fn.args();
@@ -398,7 +395,7 @@ fn has_match_all_function(expr: &Arc<dyn PhysicalExpr>) -> bool {
// check if the expr is match_all function
fn is_match_all_physical(expr: &Arc<dyn PhysicalExpr>) -> bool {
if let Some(scalar_fn) = expr.as_any().downcast_ref::<ScalarFunctionExpr>() {
if let Some(scalar_fn) = expr.downcast_ref::<ScalarFunctionExpr>() {
let name = scalar_fn.name();
name.to_lowercase() == MATCH_ALL_UDF_NAME
|| name == FUZZY_MATCH_ALL_UDF_NAME
@@ -428,7 +425,7 @@ impl TreeNodeRewriter for AddFstFieldsToProjection {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, plan: Self::Node) -> Result<Transformed<Self::Node>> {
if let Some(empty_exec) = plan.as_any().downcast_ref::<NewEmptyExec>() {
if let Some(empty_exec) = plan.downcast_ref::<NewEmptyExec>() {
let schema = empty_exec.full_schema();
// used for read field from parquet file, because match_all function do not include
// the field in argument, so when we rewrite match_all function, we need to add
@@ -658,13 +655,13 @@ mod tests {
// Verify that all column indices in FilterExec match the input schema
let _ = physical_plan.apply(&mut |node: &Arc<dyn ExecutionPlan>| -> Result<TreeNodeRecursion> {
if let Some(filter) = node.as_any().downcast_ref::<FilterExec>() {
if let Some(filter) = node.downcast_ref::<FilterExec>() {
let input_schema = filter.input().schema();
let predicate = filter.predicate();
// Traverse the predicate expression tree to find all Column references
let _ = predicate.apply(&mut |expr: &Arc<dyn PhysicalExpr>| -> Result<TreeNodeRecursion> {
if let Some(column) = expr.as_any().downcast_ref::<Column>() {
if let Some(column) = expr.downcast_ref::<Column>() {
let column_name = column.name();
let column_index = column.index();
src/service/search/datafusion/optimizer/physical_optimizer/utils.rs
+10 -12
View File
@@ -38,7 +38,7 @@ pub fn is_aggregate_exec(plan: &Arc<dyn ExecutionPlan>) -> bool {
}
pub fn extract_string_literal(expr: &Arc<dyn PhysicalExpr>) -> Result<String> {
if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
if let Some(literal) = expr.downcast_ref::<Literal>() {
match literal.value() {
ScalarValue::Utf8(Some(s)) => Ok(s.clone()),
ScalarValue::Utf8View(Some(s)) => Ok(s.to_string()),
@@ -56,7 +56,7 @@ pub fn extract_string_literal(expr: &Arc<dyn PhysicalExpr>) -> Result<String> {
}
pub fn extract_column(expr: &Arc<dyn PhysicalExpr>) -> Result<Column> {
if let Some(column) = expr.as_any().downcast_ref::<Column>() {
if let Some(column) = expr.downcast_ref::<Column>() {
Ok(column.clone())
} else {
Err(DataFusionError::Internal(
@@ -66,7 +66,7 @@ pub fn extract_column(expr: &Arc<dyn PhysicalExpr>) -> Result<Column> {
}
pub fn extract_int64_literal(expr: &Arc<dyn PhysicalExpr>) -> Result<i64> {
if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
if let Some(literal) = expr.downcast_ref::<Literal>() {
match literal.value() {
ScalarValue::Int64(Some(s)) => Ok(*s),
_ => Err(DataFusionError::Internal(format!(
@@ -100,9 +100,9 @@ fn disjunction_opt(
}
pub fn is_column(expr: &Arc<dyn PhysicalExpr>) -> bool {
if expr.as_any().downcast_ref::<Column>().is_some() {
if expr.downcast_ref::<Column>().is_some() {
true
} else if let Some(expr) = expr.as_any().downcast_ref::<CastExpr>() {
} else if let Some(expr) = expr.downcast_ref::<CastExpr>() {
is_column(expr.expr())
} else {
false
@@ -110,9 +110,9 @@ pub fn is_column(expr: &Arc<dyn PhysicalExpr>) -> bool {
}
pub fn get_column_name(expr: &Arc<dyn PhysicalExpr>) -> &str {
if let Some(expr) = expr.as_any().downcast_ref::<Column>() {
if let Some(expr) = expr.downcast_ref::<Column>() {
expr.name()
} else if let Some(expr) = expr.as_any().downcast_ref::<CastExpr>() {
} else if let Some(expr) = expr.downcast_ref::<CastExpr>() {
get_column_name(expr.expr())
} else {
UNKNOWN_NAME
@@ -133,7 +133,7 @@ pub fn is_count_rows_aggregate(expr: &AggregateFunctionExpr) -> bool {
}
pub fn is_value(expr: &Arc<dyn PhysicalExpr>) -> bool {
expr.as_any().downcast_ref::<Literal>().is_some()
expr.downcast_ref::<Literal>().is_some()
}
pub fn is_only_timestamp_filter(expr: &[&Arc<dyn PhysicalExpr>]) -> bool {
@@ -141,7 +141,7 @@ pub fn is_only_timestamp_filter(expr: &[&Arc<dyn PhysicalExpr>]) -> bool {
}
fn is_timestamp_filter(expr: &Arc<dyn PhysicalExpr>) -> bool {
if let Some(expr) = expr.as_any().downcast_ref::<BinaryExpr>() {
if let Some(expr) = expr.downcast_ref::<BinaryExpr>() {
match expr.op() {
Operator::Gt | Operator::GtEq | Operator::Lt | Operator::LtEq => {
let column = if is_value(expr.left()) && is_column(expr.right()) {
@@ -286,9 +286,7 @@ mod tests {
fn test_disjunction_empty_returns_true_literal() {
let result = disjunction(vec![]);
// empty → lit(true)
let lit = result
.as_any()
.downcast_ref::<datafusion::physical_plan::expressions::Literal>();
let lit = result.downcast_ref::<datafusion::physical_plan::expressions::Literal>();
assert!(lit.is_some());
}
src/service/search/datafusion/optimizer/utils.rs
+2 -9
View File
@@ -287,15 +287,8 @@ impl TreeNodeRewriter for SortByTime {
fn generate_table_source_with_sorted_by_time(
table_source: Arc<dyn TableSource>,
) -> Arc<dyn TableSource> {
let source: &DefaultTableSource = table_source
.as_any()
.downcast_ref::<DefaultTableSource>()
.unwrap();
if let Some(table_provider) = source
.table_provider
.as_any()
.downcast_ref::<NewEmptyTable>()
{
let source: &DefaultTableSource = table_source.downcast_ref::<DefaultTableSource>().unwrap();
if let Some(table_provider) = source.table_provider.downcast_ref::<NewEmptyTable>() {
let mut new_table_provider = (*table_provider).clone();
new_table_provider.sorted_by_time = true;
let new_source = DefaultTableSource::new(Arc::new(new_table_provider));
src/service/search/datafusion/peak_memory_pool.rs
+16 -3
View File
@@ -13,9 +13,12 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
use std::{
fmt,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use datafusion::execution::memory_pool::{
@@ -51,6 +54,10 @@ impl PeakMemoryPool {
}
impl MemoryPool for PeakMemoryPool {
fn name(&self) -> &str {
"PeakMemoryPool"
}
fn register(&self, consumer: &MemoryConsumer) {
self.inner.register(consumer);
}
@@ -91,6 +98,12 @@ impl MemoryPool for PeakMemoryPool {
}
}
impl fmt::Display for PeakMemoryPool {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}({})", self.name(), self.inner)
}
}
impl Drop for PeakMemoryPool {
fn drop(&mut self) {
let peak = self.peak_memory.load(Ordering::Relaxed);
src/service/search/datafusion/plan/deduplication_exec.rs
+3 -7
View File
@@ -14,7 +14,6 @@
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{
any::Any,
sync::Arc,
task::{Context, Poll},
};
@@ -111,10 +110,6 @@ impl ExecutionPlan for DeduplicationExec {
"DeduplicationExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.cache
}
@@ -161,7 +156,7 @@ impl ExecutionPlan for DeduplicationExec {
)))
}
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
self.input.partition_statistics(partition)
}
@@ -455,7 +450,8 @@ mod tests {
#[test]
fn test_deduplication_exec_as_any() {
let exec = make_exec();
assert!(exec.as_any().downcast_ref::<DeduplicationExec>().is_some());
let exec_plan: &dyn ExecutionPlan = &exec;
assert!(exec_plan.downcast_ref::<DeduplicationExec>().is_some());
}
#[test]
src/service/search/datafusion/plan/tantivy_optimize_exec.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, collections::HashSet, sync::Arc};
use std::{collections::HashSet, sync::Arc};
use arrow::array::{
Array, Int64Array, TimestampMicrosecondArray, TimestampNanosecondArray, UInt64Array,
@@ -114,10 +114,6 @@ impl ExecutionPlan for TantivyOptimizeExec {
"TantivyOptimizeExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.cache
}
src/service/search/datafusion/storage/file_statistics_cache.rs
+142 -147
View File
@@ -25,10 +25,10 @@ use std::{
use config::metrics;
use dashmap::DashMap;
use datafusion::{
common::Statistics,
common::{Statistics, TableReference},
execution::cache::{
CacheAccessor,
cache_manager::{CachedFileMetadata, FileStatisticsCacheEntry},
CacheAccessor, TableScopedPath,
cache_manager::{self, CachedFileMetadata, FileStatisticsCacheEntry},
},
};
use object_store::{ObjectMeta, path::Path};
@@ -145,10 +145,10 @@ impl Default for FileStatisticsCache {
}
}
impl CacheAccessor<Path, CachedFileMetadata> for FileStatisticsCache {
impl CacheAccessor<TableScopedPath, CachedFileMetadata> for FileStatisticsCache {
/// Get cached metadata for file location.
fn get(&self, k: &Path) -> Option<CachedFileMetadata> {
let k = self.format_key(k);
fn get(&self, k: &TableScopedPath) -> Option<CachedFileMetadata> {
let k = self.format_key(&k.path);
match self.statistics.get(&k) {
Some(s) => {
metrics::QUERY_PARQUET_METADATA_CACHE_HITS_TOTAL
@@ -171,8 +171,8 @@ impl CacheAccessor<Path, CachedFileMetadata> for FileStatisticsCache {
}
/// Save collected file statistics
fn put(&self, k: &Path, value: CachedFileMetadata) -> Option<CachedFileMetadata> {
let k = self.format_key(k);
fn put(&self, k: &TableScopedPath, value: CachedFileMetadata) -> Option<CachedFileMetadata> {
let k = self.format_key(&k.path);
let entry_size = Self::estimate_entry_size(&k, &value.meta, &value.statistics);
let old = self.statistics.insert(
@@ -201,8 +201,8 @@ impl CacheAccessor<Path, CachedFileMetadata> for FileStatisticsCache {
old.map(|(meta, stats, _)| CachedFileMetadata::new(meta, stats, None))
}
fn remove(&self, k: &Path) -> Option<CachedFileMetadata> {
let k = self.format_key(k);
fn remove(&self, k: &TableScopedPath) -> Option<CachedFileMetadata> {
let k = self.format_key(&k.path);
self.statistics.remove(&k).map(|(_, (meta, stats, size))| {
self.current_memory
.fetch_sub(size as i64, Ordering::Relaxed);
@@ -210,8 +210,8 @@ impl CacheAccessor<Path, CachedFileMetadata> for FileStatisticsCache {
})
}
fn contains_key(&self, k: &Path) -> bool {
let k = self.format_key(k);
fn contains_key(&self, k: &TableScopedPath) -> bool {
let k = self.format_key(&k.path);
self.statistics.contains_key(&k)
}
@@ -224,17 +224,34 @@ impl CacheAccessor<Path, CachedFileMetadata> for FileStatisticsCache {
self.cacher.lock().clear();
self.current_memory.store(0, Ordering::Relaxed);
}
fn name(&self) -> String {
"FileStatisticsCache".to_string()
}
}
impl datafusion::execution::cache::cache_manager::FileStatisticsCache for FileStatisticsCache {
fn list_entries(&self) -> HashMap<Path, FileStatisticsCacheEntry> {
let mut entries = HashMap::<Path, FileStatisticsCacheEntry>::new();
impl cache_manager::FileStatisticsCache for FileStatisticsCache {
fn cache_limit(&self) -> usize {
config::get_config()
.limit
.datafusion_file_stat_cache_max_size
}
fn update_cache_limit(&self, _limit: usize) {
// No-op: this is a process-wide, long-lived global cache whose eviction
// budget is owned solely by `put()` via `datafusion_file_stat_cache_max_size`.
// DataFusion calls this once per `CacheManager::try_new` (i.e. per session /
// per query) with its own per-session limit;
}
fn list_entries(&self) -> HashMap<TableScopedPath, FileStatisticsCacheEntry> {
let mut entries = HashMap::<TableScopedPath, FileStatisticsCacheEntry>::new();
for entry in &self.statistics {
let path = Path::from(entry.key().as_str());
let path = TableScopedPath {
table: None,
path: Path::from(entry.key().as_str()),
};
let (object_meta, stats, _) = entry.value();
entries.insert(
path,
@@ -251,86 +268,102 @@ impl datafusion::execution::cache::cache_manager::FileStatisticsCache for FileSt
entries
}
fn drop_table_entries(
&self,
_table_ref: &Option<TableReference>,
) -> datafusion::error::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use chrono::DateTime;
use chrono::{DateTime, Utc};
use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use super::*;
#[test]
fn test_file_statistics_cache() {
let meta = ObjectMeta {
location: Path::from("test"),
last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00")
.unwrap()
.into(),
size: 1024,
/// Wrap a [`Path`] into the table-scoped key the cache API expects.
fn key(location: &Path) -> TableScopedPath {
TableScopedPath {
table: None,
path: location.clone(),
}
}
/// Parse an RFC3339 timestamp, panicking on malformed input (test-only).
fn ts(s: &str) -> DateTime<Utc> {
DateTime::parse_from_rfc3339(s).unwrap().into()
}
/// Minimal [`ObjectMeta`] for a location, fixed `last_modified`, no etag/version.
fn object_meta(location: &str, size: u64) -> ObjectMeta {
ObjectMeta {
location: Path::from(location),
last_modified: ts("2024-01-15T00:00:00+00:00"),
size,
e_tag: None,
version: None,
};
let cache = FileStatisticsCache::default();
assert!(cache.get(&meta.location).is_none());
}
}
let stats: Arc<Statistics> = Statistics::new_unknown(&Schema::new(vec![Field::new(
/// Unknown statistics for the given fields, ready to cache.
fn unknown_stats(fields: Vec<Field>) -> Arc<Statistics> {
Statistics::new_unknown(&Schema::new(fields)).into()
}
/// Insert an entry into the cache under its own location.
fn put(cache: &FileStatisticsCache, meta: ObjectMeta, stats: Arc<Statistics>) {
cache.put(
&key(&meta.location),
CachedFileMetadata::new(meta, stats, None),
);
}
#[test]
fn test_file_statistics_cache() {
let meta = object_meta("test", 1024);
let cache = FileStatisticsCache::default();
assert!(cache.get(&key(&meta.location)).is_none());
let stats = unknown_stats(vec![Field::new(
"test_column",
DataType::Timestamp(TimeUnit::Second, None),
false,
)]))
.into();
cache.put(
&meta.location,
CachedFileMetadata::new(meta.clone(), stats, None),
);
let cached = cache.get(&meta.location);
assert!(cached.is_some());
assert!(cached.unwrap().is_valid_for(&meta));
)]);
put(&cache, meta.clone(), stats);
// file size changed
let mut meta2 = meta.clone();
meta2.size = 2048;
let cached = cache.get(&meta2.location);
assert!(cached.is_some());
assert!(!cached.unwrap().is_valid_for(&meta2));
// exact match is valid
let cached = cache.get(&key(&meta.location)).expect("entry present");
assert!(cached.is_valid_for(&meta));
// file last_modified changed
let mut meta2 = meta.clone();
meta2.last_modified = DateTime::parse_from_rfc3339("2022-09-27T22:40:00+02:00")
.unwrap()
.into();
let cached = cache.get(&meta2.location);
assert!(cached.is_some());
assert!(!cached.unwrap().is_valid_for(&meta2));
// same location but file size changed -> cached but stale
let mut changed = meta.clone();
changed.size = 2048;
let cached = cache.get(&key(&changed.location)).expect("entry present");
assert!(!cached.is_valid_for(&changed));
// different file
let mut meta2 = meta;
meta2.location = Path::from("test2");
assert!(cache.get(&meta2.location).is_none());
// same location but last_modified changed -> cached but stale
let mut changed = meta.clone();
changed.last_modified = ts("2024-01-15T01:00:00+00:00");
let cached = cache.get(&key(&changed.location)).expect("entry present");
assert!(!cached.is_valid_for(&changed));
// different location -> miss
assert!(cache.get(&key(&Path::from("test2"))).is_none());
}
#[test]
fn test_memory_size_calculation() {
let cache = FileStatisticsCache::new();
assert_eq!(cache.memory_size(), 0, "empty cache tracks zero memory");
// Empty cache tracks zero memory
let empty_size = cache.memory_size();
assert_eq!(empty_size, 0);
// Add some entries
for i in 0..10 {
let meta = ObjectMeta {
location: Path::from(format!("test_file_{i}")),
last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00")
.unwrap()
.into(),
size: 1024 * (i + 1),
e_tag: Some(format!("etag_{i}")),
version: Some(format!("v{i}")),
};
let schema = Schema::new(vec![
for i in 0..10u64 {
let mut meta = object_meta(&format!("test_file_{i}"), 1024 * (i + 1));
meta.e_tag = Some(format!("etag_{i}"));
meta.version = Some(format!("v{i}"));
let stats = unknown_stats(vec![
Field::new("column1", DataType::Utf8, false),
Field::new("column2", DataType::Int64, true),
Field::new(
@@ -339,84 +372,55 @@ mod tests {
false,
),
]);
cache.put(
&meta.location,
CachedFileMetadata::new(
meta.clone(),
Statistics::new_unknown(&schema).into(),
None,
),
);
put(&cache, meta, stats);
}
// Size should increase after adding entries
let filled_size = cache.memory_size();
assert!(filled_size > empty_size);
println!(
"Cache with {} entries uses {} bytes",
cache.len(),
filled_size,
assert!(
cache.memory_size() > 0,
"memory tracked after inserting entries"
);
assert_eq!(cache.len(), 10);
// After clearing, tracked memory must return to zero
cache.clear();
assert_eq!(cache.memory_size(), 0);
assert_eq!(cache.memory_size(), 0, "cleared cache returns to zero");
}
#[test]
fn test_cache_name() {
let cache = FileStatisticsCache::new();
assert_eq!(cache.name(), "FileStatisticsCache");
assert_eq!(FileStatisticsCache::new().name(), "FileStatisticsCache");
}
#[test]
fn test_cache_contains_key_and_remove() {
let cache = FileStatisticsCache::new();
let path = Path::from("test_file");
let meta = ObjectMeta {
location: path.clone(),
last_modified: DateTime::parse_from_rfc3339("2024-01-15T00:00:00+00:00")
.unwrap()
.into(),
size: 512,
e_tag: None,
version: None,
};
let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]);
let stats: Arc<Statistics> = Statistics::new_unknown(&schema).into();
let meta = object_meta("test_file", 512);
let k = key(&meta.location);
assert!(!cache.contains_key(&path));
assert!(!cache.contains_key(&k));
cache.put(&path, CachedFileMetadata::new(meta.clone(), stats, None));
assert!(cache.contains_key(&path));
put(
&cache,
meta,
unknown_stats(vec![Field::new("col", DataType::Utf8, false)]),
);
assert!(cache.contains_key(&k));
assert_eq!(cache.len(), 1);
let removed = cache.remove(&path);
assert!(removed.is_some());
assert!(!cache.contains_key(&path));
assert!(cache.remove(&k).is_some());
assert!(!cache.contains_key(&k));
assert_eq!(cache.len(), 0);
}
#[test]
fn test_cache_clear() {
let cache = FileStatisticsCache::new();
let schema = Schema::new(vec![Field::new("col", DataType::Int64, false)]);
for i in 0..3 {
let path = Path::from(format!("file_{i}"));
let meta = ObjectMeta {
location: path.clone(),
last_modified: DateTime::parse_from_rfc3339("2024-01-15T00:00:00+00:00")
.unwrap()
.into(),
size: 100 * (i + 1),
e_tag: None,
version: None,
};
let stats: Arc<Statistics> = Statistics::new_unknown(&schema).into();
cache.put(&path, CachedFileMetadata::new(meta, stats, None));
for i in 0..3u64 {
let meta = object_meta(&format!("file_{i}"), 100 * (i + 1));
put(
&cache,
meta,
unknown_stats(vec![Field::new("col", DataType::Int64, false)]),
);
}
assert_eq!(cache.len(), 3);
@@ -428,34 +432,25 @@ mod tests {
fn test_list_entries() {
use datafusion::execution::cache::cache_manager::FileStatisticsCache as FscTrait;
let cache = super::FileStatisticsCache::new();
let path = Path::from("list_test");
let meta = ObjectMeta {
location: path.clone(),
last_modified: DateTime::parse_from_rfc3339("2024-01-15T00:00:00+00:00")
.unwrap()
.into(),
size: 256,
e_tag: None,
version: None,
};
let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]);
let stats: Arc<Statistics> = Statistics::new_unknown(&schema).into();
let cache = FileStatisticsCache::new();
let meta = object_meta("list_test", 256);
let k = key(&meta.location);
let entries_before = FscTrait::list_entries(&cache);
assert!(entries_before.is_empty());
assert!(FscTrait::list_entries(&cache).is_empty());
cache.put(&path, CachedFileMetadata::new(meta, stats, None));
put(
&cache,
meta,
unknown_stats(vec![Field::new("col", DataType::Utf8, false)]),
);
let entries = FscTrait::list_entries(&cache);
assert_eq!(entries.len(), 1);
assert!(entries.contains_key(&path));
assert!(entries.contains_key(&k));
}
#[test]
fn test_cache_remove_nonexistent_returns_none() {
let cache = FileStatisticsCache::new();
let path = Path::from("does_not_exist");
let removed = cache.remove(&path);
assert!(removed.is_none());
assert!(cache.remove(&key(&Path::from("does_not_exist"))).is_none());
}
}
src/service/search/datafusion/table_provider/catalog.rs
-8
View File
@@ -14,7 +14,6 @@
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{
any::Any,
collections::HashMap,
sync::{Arc, RwLock},
};
@@ -40,10 +39,6 @@ impl StreamTypeProvider {
#[async_trait]
impl SchemaProvider for StreamTypeProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn table_names(&self) -> Vec<String> {
let tables = self.tables.read().unwrap();
tables.keys().cloned().collect::<Vec<_>>()
@@ -104,9 +99,6 @@ mod tests {
#[async_trait]
impl TableProvider for DummyTable {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema(&self) -> SchemaRef {
self.0.clone()
}
src/service/search/datafusion/table_provider/empty_table.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use async_trait::async_trait;
use datafusion::{
@@ -64,10 +64,6 @@ impl NewEmptyTable {
#[async_trait]
impl TableProvider for NewEmptyTable {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
src/service/search/datafusion/table_provider/enrich_table.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use async_trait::async_trait;
use config::TIMESTAMP_COL_NAME;
@@ -67,10 +67,6 @@ impl EnrichTable {
#[async_trait]
impl TableProvider for EnrichTable {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
src/service/search/datafusion/table_provider/listing_adapter.rs
+3 -8
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use arrow_schema::{SchemaRef, SortOptions};
use config::{TIMESTAMP_COL_NAME, get_config};
@@ -74,10 +74,6 @@ impl ListingTableAdapter {
#[async_trait]
impl TableProvider for ListingTableAdapter {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
Arc::clone(&self.listing_table.schema())
}
@@ -188,10 +184,9 @@ fn handler_tantivy_index(
order_by_time_desc: bool,
target_partitions: usize,
) -> Arc<dyn ExecutionPlan> {
if let Some(data_source_exec) = plan.as_any().downcast_ref::<DataSourceExec>()
if let Some(data_source_exec) = plan.downcast_ref::<DataSourceExec>()
&& let Some(config) = data_source_exec
.data_source()
.as_any()
.downcast_ref::<FileScanConfig>()
{
let mut file_groups = config.file_groups.clone();
@@ -250,7 +245,7 @@ fn handler_tantivy_index(
.into_iter()
.map(|mut file| {
if let Some(access_plan) = generate_access_plan(&file) {
file = file.with_extensions(access_plan);
file = file.with_extension(access_plan);
}
file
})
src/service/search/datafusion/table_provider/memtable.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use arrow::array::RecordBatch;
use arrow_schema::{DataType, SchemaRef, SortOptions};
@@ -120,10 +120,6 @@ mod tests {
#[async_trait]
impl TableProvider for NewMemTable {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
self.mem_table.schema()
}
src/service/search/datafusion/table_provider/uniontable.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::Arc};
use std::sync::Arc;
use arrow_schema::SchemaRef;
use async_trait::async_trait;
@@ -40,10 +40,6 @@ impl NewUnionTable {
#[async_trait]
impl TableProvider for NewUnionTable {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
src/service/search/datafusion/udaf/summary_percentile.rs
-4
View File
@@ -82,10 +82,6 @@ impl Default for SummaryPercentile {
}
impl AggregateUDFImpl for SummaryPercentile {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
SUMMARY_PERCENTILE
}
src/service/search/datafusion/udf/cast_to_timestamp_udf.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::LazyLock as Lazy};
use std::sync::LazyLock as Lazy;
use arrow_schema::TimeUnit;
use datafusion::{
@@ -47,10 +47,6 @@ impl CastToTimestampUdf {
}
impl ScalarUDFImpl for CastToTimestampUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
CAST_TO_TIMESTAMP_UDF_NAME
}
src/service/search/datafusion/udf/histogram_udf.rs
+1 -5
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::LazyLock as Lazy};
use std::sync::LazyLock as Lazy;
use arrow::datatypes::{DataType, DataType::Timestamp, TimeUnit::Microsecond};
use datafusion::{
@@ -40,10 +40,6 @@ impl HistogramUdf {
}
impl ScalarUDFImpl for HistogramUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
HISTOGRAM_UDF_NAME
}
src/service/search/datafusion/udf/match_all_udf.rs
+1 -9
View File
@@ -13,7 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{any::Any, sync::LazyLock as Lazy};
use std::sync::LazyLock as Lazy;
use arrow::datatypes::DataType;
use datafusion::{
@@ -44,10 +44,6 @@ impl MatchAllUdf {
}
impl ScalarUDFImpl for MatchAllUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
MATCH_ALL_UDF_NAME
}
@@ -87,10 +83,6 @@ impl FuzzyMatchAllUdf {
}
impl ScalarUDFImpl for FuzzyMatchAllUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
FUZZY_MATCH_ALL_UDF_NAME
}
src/service/search/datafusion/udf/regexp_matches_udf.rs
+1 -8
View File
@@ -13,10 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{
any::Any,
sync::{Arc, LazyLock as Lazy},
};
use std::sync::{Arc, LazyLock as Lazy};
use arrow::array::{Array, ArrayRef, GenericStringBuilder, ListBuilder, OffsetSizeTrait};
use arrow_schema::Field;
@@ -87,10 +84,6 @@ impl RegexpMatchesFunc {
}
impl ScalarUDFImpl for RegexpMatchesFunc {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
REGEX_MATCHES_UDF_NAME
}
src/service/search/datafusion/udf/regexp_udf.rs
-4
View File
@@ -247,10 +247,6 @@ impl RegxpMatchToFields {
}
impl ScalarUDFImpl for RegxpMatchToFields {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
"regexp_match_to_fields"
}
src/service/search/datafusion/udf/str_match_udf.rs
+1 -12
View File
@@ -13,10 +13,7 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::{
any::Any,
sync::{Arc, LazyLock as Lazy},
};
use std::sync::{Arc, LazyLock as Lazy};
use datafusion::{
arrow::{
@@ -54,10 +51,6 @@ impl StrMatchUdf {
}
impl ScalarUDFImpl for StrMatchUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
super::STR_MATCH_UDF_NAME
}
@@ -95,10 +88,6 @@ impl StrMatchIgnoreCaseUdf {
}
impl ScalarUDFImpl for StrMatchIgnoreCaseUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
super::STR_MATCH_UDF_IGNORE_CASE_NAME
}
src/service/search/index.rs
+10 -10
View File
@@ -267,7 +267,7 @@ impl Condition {
}
pub fn from_physical_expr(expr: &Arc<dyn PhysicalExpr>) -> Self {
if let Some(expr) = expr.as_any().downcast_ref::<BinaryExpr>() {
if let Some(expr) = expr.downcast_ref::<BinaryExpr>() {
match expr.op() {
Operator::Eq | Operator::NotEq => {
let (field, value) = if is_physical_value(expr.left())
@@ -302,11 +302,11 @@ impl Condition {
),
_ => unreachable!(),
}
} else if let Some(expr) = expr.as_any().downcast_ref::<InListExpr>() {
} else if let Some(expr) = expr.downcast_ref::<InListExpr>() {
let field = get_physical_column_name(expr.expr()).to_string();
let values = expr.list().iter().map(get_physical_value).collect();
Condition::In(field, values, expr.negated())
} else if let Some(expr) = expr.as_any().downcast_ref::<ScalarFunctionExpr>() {
} else if let Some(expr) = expr.downcast_ref::<ScalarFunctionExpr>() {
let name = expr.name();
match name {
MATCH_ALL_UDF_NAME => Condition::MatchAll(get_physical_value(&expr.args()[0])),
@@ -327,7 +327,7 @@ impl Condition {
}
_ => unreachable!(),
}
} else if let Some(expr) = expr.as_any().downcast_ref::<NotExpr>() {
} else if let Some(expr) = expr.downcast_ref::<NotExpr>() {
Condition::Not(Box::new(Condition::from_physical_expr(expr.arg())))
} else {
unreachable!()
@@ -705,9 +705,9 @@ impl Condition {
// TODO: duplication with datafusion/optimizer/physical_optimizer/utils.rs
fn is_physical_column(expr: &Arc<dyn PhysicalExpr>) -> bool {
if expr.as_any().downcast_ref::<Column>().is_some() {
if expr.downcast_ref::<Column>().is_some() {
true
} else if let Some(expr) = expr.as_any().downcast_ref::<CastExpr>() {
} else if let Some(expr) = expr.downcast_ref::<CastExpr>() {
is_physical_column(expr.expr())
} else {
false
@@ -716,9 +716,9 @@ fn is_physical_column(expr: &Arc<dyn PhysicalExpr>) -> bool {
// TODO: duplication with datafusion/optimizer/physical_optimizer/utils.rs
fn get_physical_column_name(expr: &Arc<dyn PhysicalExpr>) -> &str {
if let Some(expr) = expr.as_any().downcast_ref::<Column>() {
if let Some(expr) = expr.downcast_ref::<Column>() {
expr.name()
} else if let Some(expr) = expr.as_any().downcast_ref::<CastExpr>() {
} else if let Some(expr) = expr.downcast_ref::<CastExpr>() {
get_physical_column_name(expr.expr())
} else {
UNKNOWN_NAME
@@ -726,11 +726,11 @@ fn get_physical_column_name(expr: &Arc<dyn PhysicalExpr>) -> &str {
}
fn is_physical_value(expr: &Arc<dyn PhysicalExpr>) -> bool {
expr.as_any().downcast_ref::<Literal>().is_some()
expr.downcast_ref::<Literal>().is_some()
}
fn get_physical_value(expr: &Arc<dyn PhysicalExpr>) -> String {
if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
if let Some(literal) = expr.downcast_ref::<Literal>() {
match literal.value() {
ScalarValue::Boolean(Some(b)) => b.to_string(),
ScalarValue::Int64(Some(i)) => i.to_string(),
src/service/search/sql/rewriter/track_total_hits.rs
+2 -2
View File
@@ -89,7 +89,7 @@ impl VisitorMut for TrackTotalHitsVisitor {
value_table_mode: None,
exclude: None,
flavor: SelectFlavor::Standard,
optimizer_hint: None,
optimizer_hints: vec![],
select_modifiers: None,
})));
*query = Query {
@@ -177,7 +177,7 @@ impl VisitorMut for TrackTotalHitsVisitor {
value_table_mode: None,
exclude: None,
flavor: SelectFlavor::Standard,
optimizer_hint: None,
optimizer_hints: vec![],
select_modifiers: None,
})));
*query = Query {
src/service/search/sql/visitor/group_by.rs
+1 -1
View File
@@ -73,7 +73,7 @@ impl<'n> TreeNodeVisitor<'n> for GroupByFieldVisitor {
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
let name = node.name();
if name == "AggregateExec" {
let aggregate = node.as_any().downcast_ref::<AggregateExec>().unwrap();
let aggregate = node.downcast_ref::<AggregateExec>().unwrap();
let group_by = aggregate.group_expr();
let fields = group_by
.expr()
src/service/search/utils.rs
+1 -1
View File
@@ -78,7 +78,7 @@ impl ExecutionPlanVisitor for ScanStatsVisitor {
type Error = datafusion::common::DataFusionError;
fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> {
let mayby_remote_scan_exec = plan.as_any().downcast_ref::<RemoteScanExec>();
let mayby_remote_scan_exec = plan.downcast_ref::<RemoteScanExec>();
if let Some(remote_scan_exec) = mayby_remote_scan_exec {
{
let guard = remote_scan_exec.scan_stats.lock();
src/service/tantivy/reader/vortex.rs
+3 -2
View File
@@ -49,9 +49,10 @@ pub(super) async fn scan_vortex_async(
let stream = scan.into_array_stream()?.then(move |result| {
let data_type = data_type.clone();
let session = session.clone();
async move {
match result {
Ok(array) => vortex_array_to_record_batch(array, &data_type),
Ok(array) => vortex_array_to_record_batch(&session, array, &data_type),
Err(e) => Err(ArrowError::ExternalError(Box::new(e))),
}
}
@@ -78,7 +79,7 @@ pub(super) fn scan_vortex_row_range(
let iter: Box<dyn Iterator<Item = Result<RecordBatch, ArrowError>> + 'static> = Box::new(
scan.into_array_iter(&runtime)?
.map(move |result| match result {
Ok(array) => vortex_array_to_record_batch(array, &data_type),
Ok(array) => vortex_array_to_record_batch(&session, array, &data_type),
Err(e) => Err(ArrowError::ExternalError(Box::new(e))),
}),
);
tests/api-testing/tests/vortex/test_schema_evolution.py
+7 -3
View File
@@ -61,9 +61,13 @@ class TestNewFieldAddition:
ingest(client, s, batch2)
flush_and_wait(client, s, expected=10)
for field in ("field_a", "field_b"):
n = count_records(client, s, where=f"{field} IS NOT NULL AND {field} != ''")
assert n == 10, f"{field} should appear in all 10 rows, got {n}"
# field_a is a string — IS NOT NULL AND != '' is a valid emptiness check.
n = count_records(client, s, where="field_a IS NOT NULL AND field_a != ''")
assert n == 10, f"field_a should appear in all 10 rows, got {n}"
# field_b is int — comparing to '' is a numeric/string type mismatch; use
# the actual value (1 in every row) instead.
n = count_records(client, s, where="field_b = 1")
assert n == 10, f"field_b should appear in all 10 rows, got {n}"
# Use value-specific checks: IS NOT NULL AND != '' is unreliable for
# numeric/boolean fields (e.g. field_d=0, field_e=false may look empty).