diff --git a/src/handler/http/request/status/mod.rs b/src/handler/http/request/status/mod.rs index d0be77d6bc..e44725786c 100644 --- a/src/handler/http/request/status/mod.rs +++ b/src/handler/http/request/status/mod.rs @@ -70,7 +70,7 @@ use crate::{ }, service::{ db, - search::datafusion::{storage::file_statistics_cache, DEFAULT_FUNCTIONS}, + search::datafusion::{storage::file_statistics_cache, udf::DEFAULT_FUNCTIONS}, }, }; diff --git a/src/service/promql/engine.rs b/src/service/promql/engine.rs index a506c76b78..37de7a928a 100644 --- a/src/service/promql/engine.rs +++ b/src/service/promql/engine.rs @@ -1026,14 +1026,15 @@ async fn selector_load_data_from_datafusion( } MatchOp::Re(_re) => { let regexp_match_udf = - crate::service::search::datafusion::regexp_udf::REGEX_MATCH_UDF.clone(); + crate::service::search::datafusion::udf::regexp_udf::REGEX_MATCH_UDF.clone(); df_group = df_group.filter( regexp_match_udf.call(vec![col(mat.name.clone()), lit(mat.value.clone())]), )? } MatchOp::NotRe(_re) => { let regexp_not_match_udf = - crate::service::search::datafusion::regexp_udf::REGEX_NOT_MATCH_UDF.clone(); + crate::service::search::datafusion::udf::regexp_udf::REGEX_NOT_MATCH_UDF + .clone(); df_group = df_group.filter( regexp_not_match_udf.call(vec![col(mat.name.clone()), lit(mat.value.clone())]), )? diff --git a/src/service/search/datafusion/exec.rs b/src/service/search/datafusion/exec.rs index bcfc8de662..e0e74edb1a 100644 --- a/src/service/search/datafusion/exec.rs +++ b/src/service/search/datafusion/exec.rs @@ -37,7 +37,7 @@ use datafusion::{ common::{FileType, GetExt}, datasource::{ file_format::{json::JsonFormat, parquet::ParquetFormat}, - listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, + listing::{ListingOptions, ListingTableConfig, ListingTableUrl}, object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry}, MemTable, }, @@ -60,7 +60,9 @@ use once_cell::sync::Lazy; use parquet::arrow::ArrowWriter; use regex::Regex; -use super::{storage::file_list, transform_udf::get_all_transform}; +use super::{ + storage::file_list, table_provider::NewListingTable, udf::transform_udf::get_all_transform, +}; use crate::{ common::meta::functions::VRLResultResolver, service::search::{datafusion::rewrite, sql::Sql, RE_SELECT_WILDCARD}, @@ -615,7 +617,7 @@ pub async fn merge( } }; - let table = ListingTable::try_new(config)?; + let table = NewListingTable::try_new(config)?; ctx.register_table("tbl", Arc::new(table))?; // register UDF @@ -982,7 +984,7 @@ pub async fn convert_parquet_file( .with_listing_options(listing_options) .with_schema(schema.clone()); - let table = ListingTable::try_new(config)?; + let table = NewListingTable::try_new(config)?; ctx.register_table("tbl", Arc::new(table))?; // get all sorted data @@ -1088,7 +1090,7 @@ pub async fn merge_parquet_files( let config = ListingTableConfig::new(prefix) .with_listing_options(listing_options) .with_schema(schema.clone()); - let table = Arc::new(ListingTable::try_new(config)?); + let table = Arc::new(NewListingTable::try_new(config)?); ctx.register_table("tbl", table.clone())?; let df = ctx.sql(&query_sql).await?; @@ -1209,23 +1211,23 @@ pub async fn prepare_datafusion_context( } async fn register_udf(ctx: &mut SessionContext, _org_id: &str) { - ctx.register_udf(super::match_udf::MATCH_UDF.clone()); - ctx.register_udf(super::match_udf::MATCH_IGNORE_CASE_UDF.clone()); - ctx.register_udf(super::regexp_udf::REGEX_MATCH_UDF.clone()); - ctx.register_udf(super::regexp_udf::REGEX_NOT_MATCH_UDF.clone()); - ctx.register_udf(super::regexp_udf::REGEXP_MATCH_TO_FIELDS_UDF.clone()); - ctx.register_udf(super::time_range_udf::TIME_RANGE_UDF.clone()); - ctx.register_udf(super::date_format_udf::DATE_FORMAT_UDF.clone()); - ctx.register_udf(super::string_to_array_v2_udf::STRING_TO_ARRAY_V2_UDF.clone()); - ctx.register_udf(super::arrzip_udf::ARR_ZIP_UDF.clone()); - ctx.register_udf(super::arrindex_udf::ARR_INDEX_UDF.clone()); - ctx.register_udf(super::arr_descending_udf::ARR_DESCENDING_UDF.clone()); - ctx.register_udf(super::arrjoin_udf::ARR_JOIN_UDF.clone()); - ctx.register_udf(super::arrcount_udf::ARR_COUNT_UDF.clone()); - ctx.register_udf(super::arrsort_udf::ARR_SORT_UDF.clone()); - ctx.register_udf(super::cast_to_arr_udf::CAST_TO_ARR_UDF.clone()); - ctx.register_udf(super::spath_udf::SPATH_UDF.clone()); - ctx.register_udf(super::to_arr_string_udf::TO_ARR_STRING.clone()); + ctx.register_udf(super::udf::match_udf::MATCH_UDF.clone()); + ctx.register_udf(super::udf::match_udf::MATCH_IGNORE_CASE_UDF.clone()); + ctx.register_udf(super::udf::regexp_udf::REGEX_MATCH_UDF.clone()); + ctx.register_udf(super::udf::regexp_udf::REGEX_NOT_MATCH_UDF.clone()); + ctx.register_udf(super::udf::regexp_udf::REGEXP_MATCH_TO_FIELDS_UDF.clone()); + ctx.register_udf(super::udf::time_range_udf::TIME_RANGE_UDF.clone()); + ctx.register_udf(super::udf::date_format_udf::DATE_FORMAT_UDF.clone()); + ctx.register_udf(super::udf::string_to_array_v2_udf::STRING_TO_ARRAY_V2_UDF.clone()); + ctx.register_udf(super::udf::arrzip_udf::ARR_ZIP_UDF.clone()); + ctx.register_udf(super::udf::arrindex_udf::ARR_INDEX_UDF.clone()); + ctx.register_udf(super::udf::arr_descending_udf::ARR_DESCENDING_UDF.clone()); + ctx.register_udf(super::udf::arrjoin_udf::ARR_JOIN_UDF.clone()); + ctx.register_udf(super::udf::arrcount_udf::ARR_COUNT_UDF.clone()); + ctx.register_udf(super::udf::arrsort_udf::ARR_SORT_UDF.clone()); + ctx.register_udf(super::udf::cast_to_arr_udf::CAST_TO_ARR_UDF.clone()); + ctx.register_udf(super::udf::spath_udf::SPATH_UDF.clone()); + ctx.register_udf(super::udf::to_arr_string_udf::TO_ARR_STRING.clone()); { let udf_list = get_all_transform(_org_id).await; @@ -1306,7 +1308,7 @@ pub async fn register_table( } else { config = config.with_schema(schema); } - let mut table = ListingTable::try_new(config)?; + let mut table = NewListingTable::try_new(config)?; if session.storage_type != StorageType::Tmpfs { table = table.with_cache(ctx.runtime_env().cache_manager.get_file_statistic_cache()); } diff --git a/src/service/search/datafusion/mod.rs b/src/service/search/datafusion/mod.rs index ca936ce32c..28de23dfb8 100644 --- a/src/service/search/datafusion/mod.rs +++ b/src/service/search/datafusion/mod.rs @@ -15,28 +15,11 @@ use std::str::FromStr; -use config::utils::json; - -use crate::common::meta::functions::ZoFunction; - -pub mod arr_descending_udf; -pub mod arrcount_udf; -pub mod arrindex_udf; -pub mod arrjoin_udf; -pub mod arrsort_udf; -pub mod arrzip_udf; -pub mod cast_to_arr_udf; -mod date_format_udf; pub mod exec; -pub mod match_udf; -pub mod regexp_udf; pub mod rewrite; -pub mod spath_udf; pub mod storage; -pub mod string_to_array_v2_udf; -mod time_range_udf; -pub mod to_arr_string_udf; -mod transform_udf; +pub mod table_provider; +pub mod udf; #[derive(PartialEq, Debug)] pub enum MemoryPoolType { @@ -57,55 +40,3 @@ impl FromStr for MemoryPoolType { } } } - -/// The name of the match UDF given to DataFusion. -pub const MATCH_UDF_NAME: &str = "str_match"; -/// The name of the match_ignore_case UDF given to DataFusion. -pub const MATCH_UDF_IGNORE_CASE_NAME: &str = "str_match_ignore_case"; -/// The name of the regex_match UDF given to DataFusion. -pub const REGEX_MATCH_UDF_NAME: &str = "re_match"; -/// The name of the not_regex_match UDF given to DataFusion. -pub const REGEX_NOT_MATCH_UDF_NAME: &str = "re_not_match"; - -pub fn stringify_json_value(field: &json::Value) -> String { - match field { - serde_json::Value::Bool(b) => b.to_string(), - serde_json::Value::Number(n) => match n.as_f64() { - Some(f) => f.to_string(), - None => n.as_i64().unwrap().to_string(), - }, - serde_json::Value::String(s) => s.clone(), - _ => json::to_string(field).expect("failed to stringify json field"), - } -} - -pub const DEFAULT_FUNCTIONS: [ZoFunction; 7] = [ - ZoFunction { - name: "match_all_raw", - text: "match_all_raw('v')", - }, - ZoFunction { - name: "match_all_raw_ignore_case", - text: "match_all_raw_ignore_case('v')", - }, - ZoFunction { - name: "match_all", - text: "match_all('v')", - }, - ZoFunction { - name: MATCH_UDF_NAME, - text: "match_all('v')", - }, - ZoFunction { - name: MATCH_UDF_IGNORE_CASE_NAME, - text: "match_all_ignore_case('v')", - }, - ZoFunction { - name: REGEX_MATCH_UDF_NAME, - text: "re_match(field, 'pattern')", - }, - ZoFunction { - name: REGEX_NOT_MATCH_UDF_NAME, - text: "re_not_match(field, 'pattern')", - }, -]; diff --git a/src/service/search/datafusion/table_provider/helpers.rs b/src/service/search/datafusion/table_provider/helpers.rs new file mode 100644 index 0000000000..36194c728f --- /dev/null +++ b/src/service/search/datafusion/table_provider/helpers.rs @@ -0,0 +1,887 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Helper functions for the table implementation + +use std::{collections::HashMap, sync::Arc}; + +use arrow::{ + array::{Array, ArrayRef, AsArray, StringBuilder}, + compute::{and, cast, prep_null_mask_filter}, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, +}; +use arrow_schema::Fields; +use datafusion::{ + common::{ + tree_node::{TreeNode, TreeNodeRecursion}, + Column, DFSchema, DataFusionError, Result, + }, + datasource::listing::{ListingTableUrl, PartitionedFile}, + execution::context::SessionState, + physical_expr::create_physical_expr, + scalar::ScalarValue, +}; +use datafusion_expr::{execution_props::ExecutionProps, BinaryExpr, Expr, Operator, Volatility}; +use futures::{ + stream::{BoxStream, FuturesUnordered}, + StreamExt, TryStreamExt, +}; +use log::{debug, trace}; +use object_store::{ + path::{Path, DELIMITER}, + ObjectMeta, ObjectStore, +}; + +/// The maximum number of concurrent listing requests +const CONCURRENCY_LIMIT: usize = 100; + +/// Check whether the given expression can be resolved using only the columns `col_names`. +/// This means that if this function returns true: +/// - the table provider can filter the table partition values with this expression +/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering +/// was performed +pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool { + let mut is_applicable = true; + expr.apply(|expr| { + match expr { + Expr::Column(Column { ref name, .. }) => { + is_applicable &= col_names.contains(name); + if is_applicable { + Ok(TreeNodeRecursion::Jump) + } else { + Ok(TreeNodeRecursion::Stop) + } + } + Expr::Literal(_) + | Expr::Alias(_) + | Expr::OuterReferenceColumn(..) + | Expr::ScalarVariable(..) + | Expr::Not(_) + | Expr::IsNotNull(_) + | Expr::IsNull(_) + | Expr::IsTrue(_) + | Expr::IsFalse(_) + | Expr::IsUnknown(_) + | Expr::IsNotTrue(_) + | Expr::IsNotFalse(_) + | Expr::IsNotUnknown(_) + | Expr::Negative(_) + | Expr::Cast { .. } + | Expr::TryCast { .. } + | Expr::BinaryExpr { .. } + | Expr::Between { .. } + | Expr::Like { .. } + | Expr::SimilarTo { .. } + | Expr::InList { .. } + | Expr::Exists { .. } + | Expr::InSubquery(_) + | Expr::ScalarSubquery(_) + | Expr::GroupingSet(_) + | Expr::Case { .. } => Ok(TreeNodeRecursion::Continue), + + Expr::ScalarFunction(scalar_function) => { + match scalar_function.func.signature().volatility { + Volatility::Immutable => Ok(TreeNodeRecursion::Continue), + // TODO: Stable functions could be `applicable`, but that would require access + // to the context + Volatility::Stable | Volatility::Volatile => { + is_applicable = false; + Ok(TreeNodeRecursion::Stop) + } + } + } + + // TODO other expressions are not handled yet: + // - AGGREGATE, WINDOW and SORT should not end up in filter conditions, except maybe in + // some edge cases + // - Can `Wildcard` be considered as a `Literal`? + // - ScalarVariable could be `applicable`, but that would require access to the context + Expr::AggregateFunction { .. } + | Expr::Sort { .. } + | Expr::WindowFunction { .. } + | Expr::Wildcard { .. } + | Expr::Unnest { .. } + | Expr::Placeholder(_) => { + is_applicable = false; + Ok(TreeNodeRecursion::Stop) + } + } + }) + .unwrap(); + is_applicable +} + +/// Partition the list of files into `n` groups +pub fn split_files( + mut partitioned_files: Vec, + n: usize, +) -> Vec> { + if partitioned_files.is_empty() { + return vec![]; + } + + // ObjectStore::list does not guarantee any consistent order and for some + // implementations such as LocalFileSystem, it may be inconsistent. Thus + // Sort files by path to ensure consistent plans when run more than once. + partitioned_files.sort_by(|a, b| a.path().cmp(b.path())); + + // effectively this is div with rounding up instead of truncating + let chunk_size = (partitioned_files.len() + n - 1) / n; + partitioned_files + .chunks(chunk_size) + .map(|c| c.to_vec()) + .collect() +} + +struct Partition { + /// The path to the partition, including the table prefix + path: Path, + /// How many path segments below the table prefix `path` contains + /// or equivalently the number of partition values in `path` + depth: usize, + /// The files contained as direct children of this `Partition` if known + files: Option>, +} + +impl Partition { + /// List the direct children of this partition updating `self.files` with + /// any child files, and returning a list of child "directories" + async fn list(mut self, store: &dyn ObjectStore) -> Result<(Self, Vec)> { + trace!("Listing partition {}", self.path); + let prefix = Some(&self.path).filter(|p| !p.as_ref().is_empty()); + let result = store.list_with_delimiter(prefix).await?; + self.files = Some(result.objects); + Ok((self, result.common_prefixes)) + } +} + +/// Returns a recursive list of the partitions in `table_path` up to `max_depth` +async fn list_partitions( + store: &dyn ObjectStore, + table_path: &ListingTableUrl, + max_depth: usize, + partition_prefix: Option, +) -> Result> { + let partition = Partition { + path: match partition_prefix { + Some(prefix) => Path::from_iter( + Path::from(table_path.prefix().as_ref()) + .parts() + .chain(Path::from(prefix.as_ref()).parts()), + ), + None => table_path.prefix().clone(), + }, + depth: 0, + files: None, + }; + + let mut out = Vec::with_capacity(64); + + let mut pending = vec![]; + let mut futures = FuturesUnordered::new(); + futures.push(partition.list(store)); + + while let Some((partition, paths)) = futures.next().await.transpose()? { + // If pending contains a future it implies prior to this iteration + // `futures.len == CONCURRENCY_LIMIT`. We can therefore add a single + // future from `pending` to the working set + if let Some(next) = pending.pop() { + futures.push(next) + } + + let depth = partition.depth; + out.push(partition); + for path in paths { + let child = Partition { + path, + depth: depth + 1, + files: None, + }; + match depth < max_depth { + true => match futures.len() < CONCURRENCY_LIMIT { + true => futures.push(child.list(store)), + false => pending.push(child.list(store)), + }, + false => out.push(child), + } + } + } + Ok(out) +} + +async fn prune_partitions( + table_path: &ListingTableUrl, + partitions: Vec, + filters: &[Expr], + partition_cols: &[(String, DataType)], +) -> Result> { + if filters.is_empty() { + return Ok(partitions); + } + + let mut builders: Vec<_> = (0..partition_cols.len()) + .map(|_| StringBuilder::with_capacity(partitions.len(), partitions.len() * 10)) + .collect(); + + for partition in &partitions { + let cols = partition_cols.iter().map(|x| x.0.as_str()); + let parsed = + parse_partitions_for_path(table_path, &partition.path, cols).unwrap_or_default(); + + let mut builders = builders.iter_mut(); + for (p, b) in parsed.iter().zip(&mut builders) { + b.append_value(p); + } + builders.for_each(|b| b.append_null()); + } + + let arrays = partition_cols + .iter() + .zip(builders) + .map(|((_, d), mut builder)| { + let array = builder.finish(); + cast(&array, d) + }) + .collect::>()?; + + let fields: Fields = partition_cols + .iter() + .map(|(n, d)| Field::new(n, d.clone(), true)) + .collect(); + let schema = Arc::new(Schema::new(fields)); + + let df_schema = DFSchema::from_unqualifed_fields( + partition_cols + .iter() + .map(|(n, d)| Field::new(n, d.clone(), true)) + .collect(), + Default::default(), + )?; + + let batch = RecordBatch::try_new(schema.clone(), arrays)?; + + // TODO: Plumb this down + let props = ExecutionProps::new(); + + // Applies `filter` to `batch` returning `None` on error + let do_filter = |filter| -> Option { + let expr = create_physical_expr(filter, &df_schema, &props).ok()?; + expr.evaluate(&batch) + .ok()? + .into_array(partitions.len()) + .ok() + }; + + //.Compute the conjunction of the filters, ignoring errors + let mask = filters + .iter() + .fold(None, |acc, filter| match (acc, do_filter(filter)) { + (Some(a), Some(b)) => Some(and(&a, b.as_boolean()).unwrap_or(a)), + (None, Some(r)) => Some(r.as_boolean().clone()), + (r, None) => r, + }); + + let mask = match mask { + Some(mask) => mask, + None => return Ok(partitions), + }; + + // Don't retain partitions that evaluated to null + let prepared = match mask.null_count() { + 0 => mask, + _ => prep_null_mask_filter(&mask), + }; + + // Sanity check + assert_eq!(prepared.len(), partitions.len()); + + let filtered = partitions + .into_iter() + .zip(prepared.values()) + .filter_map(|(p, f)| f.then_some(p)) + .collect(); + + Ok(filtered) +} + +#[derive(Debug)] +enum PartitionValue { + Single(String), + Multi, +} + +fn populate_partition_values<'a>( + partition_values: &mut HashMap<&'a str, PartitionValue>, + filter: &'a Expr, +) { + if let Expr::BinaryExpr(BinaryExpr { + ref left, + op, + ref right, + }) = filter + { + match op { + Operator::Eq => match (left.as_ref(), right.as_ref()) { + (Expr::Column(Column { ref name, .. }), Expr::Literal(val)) + | (Expr::Literal(val), Expr::Column(Column { ref name, .. })) => { + if partition_values + .insert(name, PartitionValue::Single(val.to_string())) + .is_some() + { + partition_values.insert(name, PartitionValue::Multi); + } + } + _ => {} + }, + Operator::And => { + populate_partition_values(partition_values, left); + populate_partition_values(partition_values, right); + } + _ => {} + } + } +} + +fn evaluate_partition_prefix<'a>( + partition_cols: &'a [(String, DataType)], + filters: &'a [Expr], +) -> Option { + let mut partition_values = HashMap::new(); + for filter in filters { + populate_partition_values(&mut partition_values, filter); + } + + if partition_values.is_empty() { + return None; + } + + let mut parts = vec![]; + for (p, _) in partition_cols { + match partition_values.get(p.as_str()) { + Some(PartitionValue::Single(val)) => { + // if a partition only has a single literal value, then it can be added to the + // prefix + parts.push(format!("{p}={val}")); + } + _ => { + // break on the first unconstrainted partition to create a common prefix + // for all covered partitions. + break; + } + } + } + + if parts.is_empty() { + None + } else { + Some(Path::from_iter(parts)) + } +} + +/// Discover the partitions on the given path and prune out files +/// that belong to irrelevant partitions using `filters` expressions. +/// `filters` might contain expressions that can be resolved only at the +/// file level (e.g. Parquet row group pruning). +pub async fn pruned_partition_list<'a>( + ctx: &'a SessionState, + store: &'a dyn ObjectStore, + table_path: &'a ListingTableUrl, + filters: &'a [Expr], + file_extension: &'a str, + partition_cols: &'a [(String, DataType)], +) -> Result>> { + // if no partition col => simply list all the files + if partition_cols.is_empty() { + return Ok(Box::pin( + table_path + .list_all_files(ctx, store, file_extension) + .await? + .map_ok(|object_meta| object_meta.into()), + )); + } + + let partition_prefix = evaluate_partition_prefix(partition_cols, filters); + let partitions = + list_partitions(store, table_path, partition_cols.len(), partition_prefix).await?; + debug!("Listed {} partitions", partitions.len()); + + let pruned = prune_partitions(table_path, partitions, filters, partition_cols).await?; + + debug!("Pruning yielded {} partitions", pruned.len()); + + let stream = futures::stream::iter(pruned) + .map(move |partition: Partition| async move { + let cols = partition_cols.iter().map(|x| x.0.as_str()); + let parsed = parse_partitions_for_path(table_path, &partition.path, cols); + + let partition_values = parsed + .into_iter() + .flatten() + .zip(partition_cols) + .map(|(parsed, (_, datatype))| { + ScalarValue::try_from_string(parsed.to_string(), datatype) + }) + .collect::>>()?; + + let files = match partition.files { + Some(files) => files, + None => { + trace!("Recursively listing partition {}", partition.path); + store.list(Some(&partition.path)).try_collect().await? + } + }; + let files = files.into_iter().filter(move |o| { + let extension_match = o.location.as_ref().ends_with(file_extension); + // here need to scan subdirectories(`listing_table_ignore_subdirectory` = false) + let glob_match = table_path.contains(&o.location, false); + extension_match && glob_match + }); + + let stream = futures::stream::iter(files.map(move |object_meta| { + Ok(PartitionedFile { + object_meta, + partition_values: partition_values.clone(), + range: None, + statistics: None, + extensions: None, + }) + })); + + Ok::<_, DataFusionError>(stream) + }) + .buffer_unordered(CONCURRENCY_LIMIT) + .try_flatten() + .boxed(); + Ok(stream) +} + +/// Extract the partition values for the given `file_path` (in the given `table_path`) +/// associated to the partitions defined by `table_partition_cols` +fn parse_partitions_for_path<'a, I>( + table_path: &ListingTableUrl, + file_path: &'a Path, + table_partition_cols: I, +) -> Option> +where + I: IntoIterator, +{ + let subpath = listing_table_url_strip_prefix(table_path, file_path)?; + + let mut part_values = vec![]; + for (part, pn) in subpath.zip(table_partition_cols) { + match part.split_once('=') { + Some((name, val)) if name == pn => part_values.push(val), + _ => { + debug!( + "Ignoring file: file_path='{}', table_path='{}', part='{}', partition_col='{}'", + file_path, table_path, part, pn, + ); + return None; + } + } + } + Some(part_values) +} + +/// Strips the prefix of this [`ListingTableUrl`] from the provided path, returning +/// an iterator of the remaining path segments +pub(crate) fn listing_table_url_strip_prefix<'a, 'b: 'a>( + table_path: &'a ListingTableUrl, + path: &'b Path, +) -> Option + 'a> { + let mut stripped = path.as_ref().strip_prefix(table_path.prefix().as_ref())?; + if !stripped.is_empty() && !table_path.prefix().as_ref().is_empty() { + stripped = stripped.strip_prefix(DELIMITER)?; + } + Some(stripped.split_terminator(DELIMITER)) +} + +#[cfg(test)] +mod tests { + use std::ops::Not; + + use datafusion::logical_expr::{case, col, lit}; + + use super::*; + use crate::service::search::datafusion::table_provider::test_util::make_test_store_and_state; + + #[test] + fn test_split_files() { + let new_partitioned_file = |path: &str| PartitionedFile::new(path.to_owned(), 10); + let files = vec![ + new_partitioned_file("a"), + new_partitioned_file("b"), + new_partitioned_file("c"), + new_partitioned_file("d"), + new_partitioned_file("e"), + ]; + + let chunks = split_files(files.clone(), 1); + assert_eq!(1, chunks.len()); + assert_eq!(5, chunks[0].len()); + + let chunks = split_files(files.clone(), 2); + assert_eq!(2, chunks.len()); + assert_eq!(3, chunks[0].len()); + assert_eq!(2, chunks[1].len()); + + let chunks = split_files(files.clone(), 5); + assert_eq!(5, chunks.len()); + assert_eq!(1, chunks[0].len()); + assert_eq!(1, chunks[1].len()); + assert_eq!(1, chunks[2].len()); + assert_eq!(1, chunks[3].len()); + assert_eq!(1, chunks[4].len()); + + let chunks = split_files(files, 123); + assert_eq!(5, chunks.len()); + assert_eq!(1, chunks[0].len()); + assert_eq!(1, chunks[1].len()); + assert_eq!(1, chunks[2].len()); + assert_eq!(1, chunks[3].len()); + assert_eq!(1, chunks[4].len()); + + let chunks = split_files(vec![], 2); + assert_eq!(0, chunks.len()); + } + + #[tokio::test] + async fn test_pruned_partition_list_empty() { + let (store, state) = make_test_store_and_state(&[ + ("tablepath/mypartition=val1/notparquetfile", 100), + ("tablepath/file.parquet", 100), + ]); + let filter = Expr::eq(col("mypartition"), lit("val1")); + let pruned = pruned_partition_list( + &state, + store.as_ref(), + &ListingTableUrl::parse("file:///tablepath/").unwrap(), + &[filter], + ".parquet", + &[(String::from("mypartition"), DataType::Utf8)], + ) + .await + .expect("partition pruning failed") + .collect::>() + .await; + + assert_eq!(pruned.len(), 0); + } + + #[tokio::test] + async fn test_pruned_partition_list() { + let (store, state) = make_test_store_and_state(&[ + ("tablepath/mypartition=val1/file.parquet", 100), + ("tablepath/mypartition=val2/file.parquet", 100), + ("tablepath/mypartition=val1/other=val3/file.parquet", 100), + ]); + let filter = Expr::eq(col("mypartition"), lit("val1")); + let pruned = pruned_partition_list( + &state, + store.as_ref(), + &ListingTableUrl::parse("file:///tablepath/").unwrap(), + &[filter], + ".parquet", + &[(String::from("mypartition"), DataType::Utf8)], + ) + .await + .expect("partition pruning failed") + .try_collect::>() + .await + .unwrap(); + + assert_eq!(pruned.len(), 2); + let f1 = &pruned[0]; + assert_eq!( + f1.object_meta.location.as_ref(), + "tablepath/mypartition=val1/file.parquet" + ); + assert_eq!(&f1.partition_values, &[ScalarValue::from("val1")]); + let f2 = &pruned[1]; + assert_eq!( + f2.object_meta.location.as_ref(), + "tablepath/mypartition=val1/other=val3/file.parquet" + ); + assert_eq!(f2.partition_values, &[ScalarValue::from("val1"),]); + } + + #[tokio::test] + async fn test_pruned_partition_list_multi() { + let (store, state) = make_test_store_and_state(&[ + ("tablepath/part1=p1v1/file.parquet", 100), + ("tablepath/part1=p1v2/part2=p2v1/file1.parquet", 100), + ("tablepath/part1=p1v2/part2=p2v1/file2.parquet", 100), + ("tablepath/part1=p1v3/part2=p2v1/file2.parquet", 100), + ("tablepath/part1=p1v2/part2=p2v2/file2.parquet", 100), + ]); + let filter1 = Expr::eq(col("part1"), lit("p1v2")); + let filter2 = Expr::eq(col("part2"), lit("p2v1")); + // filter3 cannot be resolved at partition pruning + let filter3 = Expr::eq(col("part2"), col("other")); + let pruned = pruned_partition_list( + &state, + store.as_ref(), + &ListingTableUrl::parse("file:///tablepath/").unwrap(), + &[filter1, filter2, filter3], + ".parquet", + &[ + (String::from("part1"), DataType::Utf8), + (String::from("part2"), DataType::Utf8), + ], + ) + .await + .expect("partition pruning failed") + .try_collect::>() + .await + .unwrap(); + + assert_eq!(pruned.len(), 2); + let f1 = &pruned[0]; + assert_eq!( + f1.object_meta.location.as_ref(), + "tablepath/part1=p1v2/part2=p2v1/file1.parquet" + ); + assert_eq!( + &f1.partition_values, + &[ScalarValue::from("p1v2"), ScalarValue::from("p2v1"),] + ); + let f2 = &pruned[1]; + assert_eq!( + f2.object_meta.location.as_ref(), + "tablepath/part1=p1v2/part2=p2v1/file2.parquet" + ); + assert_eq!( + &f2.partition_values, + &[ScalarValue::from("p1v2"), ScalarValue::from("p2v1")] + ); + } + + #[test] + fn test_parse_partitions_for_path() { + assert_eq!( + Some(vec![]), + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), + &Path::from("bucket/mytable/file.csv"), + vec![] + ) + ); + assert_eq!( + None, + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/othertable").unwrap(), + &Path::from("bucket/mytable/file.csv"), + vec![] + ) + ); + assert_eq!( + None, + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), + &Path::from("bucket/mytable/file.csv"), + vec!["mypartition"] + ) + ); + assert_eq!( + Some(vec!["v1"]), + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), + &Path::from("bucket/mytable/mypartition=v1/file.csv"), + vec!["mypartition"] + ) + ); + assert_eq!( + Some(vec!["v1"]), + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable/").unwrap(), + &Path::from("bucket/mytable/mypartition=v1/file.csv"), + vec!["mypartition"] + ) + ); + // Only hive style partitioning supported for now: + assert_eq!( + None, + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), + &Path::from("bucket/mytable/v1/file.csv"), + vec!["mypartition"] + ) + ); + assert_eq!( + Some(vec!["v1", "v2"]), + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), + &Path::from("bucket/mytable/mypartition=v1/otherpartition=v2/file.csv"), + vec!["mypartition", "otherpartition"] + ) + ); + assert_eq!( + Some(vec!["v1"]), + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), + &Path::from("bucket/mytable/mypartition=v1/otherpartition=v2/file.csv"), + vec!["mypartition"] + ) + ); + } + + #[test] + fn test_expr_applicable_for_cols() { + assert!(expr_applicable_for_cols( + &[String::from("c1")], + &Expr::eq(col("c1"), lit("value")) + )); + assert!(!expr_applicable_for_cols( + &[String::from("c1")], + &Expr::eq(col("c2"), lit("value")) + )); + assert!(!expr_applicable_for_cols( + &[String::from("c1")], + &Expr::eq(col("c1"), col("c2")) + )); + assert!(expr_applicable_for_cols( + &[String::from("c1"), String::from("c2")], + &Expr::eq(col("c1"), col("c2")) + )); + assert!(expr_applicable_for_cols( + &[String::from("c1"), String::from("c2")], + &(Expr::eq(col("c1"), col("c2").alias("c2_alias"))).not() + )); + assert!(expr_applicable_for_cols( + &[String::from("c1"), String::from("c2")], + &(case(col("c1")) + .when(lit("v1"), lit(true)) + .otherwise(lit(false)) + .expect("valid case expr")) + )); + // static expression not relvant in this context but we + // test it as an edge case anyway in case we want to generalize + // this helper function + assert!(expr_applicable_for_cols(&[], &lit(true))); + } + + #[test] + fn test_evaluate_partition_prefix() { + let partitions = &[ + ("a".to_string(), DataType::Utf8), + ("b".to_string(), DataType::Int16), + ("c".to_string(), DataType::Boolean), + ]; + + assert_eq!( + evaluate_partition_prefix(partitions, &[col("a").eq(lit("foo"))]), + Some(Path::from("a=foo")), + ); + + assert_eq!( + evaluate_partition_prefix(partitions, &[lit("foo").eq(col("a"))]), + Some(Path::from("a=foo")), + ); + + assert_eq!( + evaluate_partition_prefix( + partitions, + &[col("a").eq(lit("foo")).and(col("b").eq(lit("bar")))], + ), + Some(Path::from("a=foo/b=bar")), + ); + + assert_eq!( + evaluate_partition_prefix( + partitions, + // list of filters should be evaluated as AND + &[col("a").eq(lit("foo")), col("b").eq(lit("bar")),], + ), + Some(Path::from("a=foo/b=bar")), + ); + + assert_eq!( + evaluate_partition_prefix( + partitions, + &[col("a") + .eq(lit("foo")) + .and(col("b").eq(lit("1"))) + .and(col("c").eq(lit("true")))], + ), + Some(Path::from("a=foo/b=1/c=true")), + ); + + // no prefix when filter is empty + assert_eq!(evaluate_partition_prefix(partitions, &[]), None); + + // b=foo results in no prefix because a is not restricted + assert_eq!( + evaluate_partition_prefix(partitions, &[Expr::eq(col("b"), lit("foo"))]), + None, + ); + + // a=foo and c=baz only results in preifx a=foo because b is not restricted + assert_eq!( + evaluate_partition_prefix( + partitions, + &[col("a").eq(lit("foo")).and(col("c").eq(lit("baz")))], + ), + Some(Path::from("a=foo")), + ); + + // partition with multiple values results in no prefix + assert_eq!( + evaluate_partition_prefix( + partitions, + &[Expr::and(col("a").eq(lit("foo")), col("a").eq(lit("bar")))], + ), + None, + ); + + // no prefix because partition a is not restricted to a single literal + assert_eq!( + evaluate_partition_prefix( + partitions, + &[Expr::or(col("a").eq(lit("foo")), col("a").eq(lit("bar")))], + ), + None, + ); + assert_eq!( + evaluate_partition_prefix(partitions, &[col("b").lt(lit(5))],), + None, + ); + } + + #[test] + fn test_evaluate_date_partition_prefix() { + let partitions = &[("a".to_string(), DataType::Date32)]; + assert_eq!( + evaluate_partition_prefix( + partitions, + &[col("a").eq(Expr::Literal(ScalarValue::Date32(Some(3))))], + ), + Some(Path::from("a=1970-01-04")), + ); + + let partitions = &[("a".to_string(), DataType::Date64)]; + assert_eq!( + evaluate_partition_prefix( + partitions, + &[col("a").eq(Expr::Literal(ScalarValue::Date64(Some( + 4 * 24 * 60 * 60 * 1000 + )))),], + ), + Some(Path::from("a=1970-01-05")), + ); + } +} diff --git a/src/service/search/datafusion/table_provider/mod.rs b/src/service/search/datafusion/table_provider/mod.rs new file mode 100644 index 0000000000..9f025d4313 --- /dev/null +++ b/src/service/search/datafusion/table_provider/mod.rs @@ -0,0 +1,379 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::{any::Any, sync::Arc}; + +use arrow_schema::{Field, Schema, SchemaBuilder, SchemaRef, SortOptions}; +use async_trait::async_trait; +use datafusion::{ + common::{plan_err, project_schema, Result, Statistics, ToDFSchema}, + datasource::{ + get_statistics_with_limit, + listing::{ListingOptions, ListingTableConfig, ListingTableUrl, PartitionedFile}, + physical_plan::FileScanConfig, + TableProvider, + }, + error::DataFusionError, + execution::{ + cache::{cache_manager::FileStatisticsCache, cache_unit::DefaultFileStatisticsCache}, + context::SessionState, + }, + physical_expr::{create_physical_expr, expressions, LexOrdering, PhysicalSortExpr}, + physical_plan::{empty::EmptyExec, ExecutionPlan}, +}; +use datafusion_expr::{utils::conjunction, Expr, TableProviderFilterPushDown, TableType}; +use futures::{future, stream, StreamExt}; +use helpers::{expr_applicable_for_cols, pruned_partition_list, split_files}; +use object_store::ObjectStore; + +mod helpers; + +#[cfg(test)] +mod test_util; + +pub(crate) struct NewListingTable { + table_paths: Vec, + /// File fields only + file_schema: SchemaRef, + /// File fields + partition columns + table_schema: SchemaRef, + options: ListingOptions, + collected_statistics: FileStatisticsCache, +} + +impl NewListingTable { + /// Create new [`ListingTable`] that lists the FS to get the files + /// to scan. See [`ListingTable`] for and example. + /// + /// Takes a `ListingTableConfig` as input which requires an `ObjectStore` and `table_path`. + /// `ListingOptions` and `SchemaRef` are optional. If they are not + /// provided the file type is inferred based on the file suffix. + /// If the schema is provided then it must be resolved before creating the table + /// and should contain the fields of the file without the table + /// partitioning columns. + pub fn try_new(config: ListingTableConfig) -> Result { + let file_schema = config + .file_schema + .ok_or_else(|| DataFusionError::Internal("No schema provided.".into()))?; + + let options = config + .options + .ok_or_else(|| DataFusionError::Internal("No ListingOptions provided".into()))?; + + // Add the partition columns to the file schema + let mut builder = SchemaBuilder::from(file_schema.as_ref().to_owned()); + for (part_col_name, part_col_type) in &options.table_partition_cols { + builder.push(Field::new(part_col_name, part_col_type.clone(), false)); + } + + let table = Self { + table_paths: config.table_paths, + file_schema, + table_schema: Arc::new(builder.finish()), + options, + collected_statistics: Arc::new(DefaultFileStatisticsCache::default()), + }; + + Ok(table) + } + + /// Set the [`FileStatisticsCache`] used to cache parquet file statistics. + /// + /// Setting a statistics cache on the `SessionContext` can avoid refetching statistics + /// multiple times in the same session. + /// + /// If `None`, creates a new [`DefaultFileStatisticsCache`] scoped to this query. + pub fn with_cache(mut self, cache: Option) -> Self { + self.collected_statistics = + cache.unwrap_or(Arc::new(DefaultFileStatisticsCache::default())); + self + } + + /// If file_sort_order is specified, creates the appropriate physical expressions + fn try_create_output_ordering(&self) -> Result> { + create_ordering(&self.table_schema, &self.options.file_sort_order) + } + + /// Get the list of files for a scan as well as the file level statistics. + /// The list is grouped to let the execution plan know how the files should + /// be distributed to different threads / executors. + async fn list_files_for_scan<'a>( + &'a self, + ctx: &'a SessionState, + filters: &'a [Expr], + limit: Option, + ) -> Result<(Vec>, Statistics)> { + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? + } else { + return Ok((vec![], Statistics::new_unknown(&self.file_schema))); + }; + // list files (with partitions) + let file_list = future::try_join_all(self.table_paths.iter().map(|table_path| { + pruned_partition_list( + ctx, + store.as_ref(), + table_path, + filters, + &self.options.file_extension, + &self.options.table_partition_cols, + ) + })) + .await?; + let file_list = stream::iter(file_list).flatten(); + // collect the statistics if required by the config + let files = file_list + .map(|part_file| async { + let mut part_file = part_file?; + if self.options.collect_stat { + let statistics = self.do_collect_statistics(ctx, &store, &part_file).await?; + part_file.statistics = Some(statistics.clone()); + Ok((part_file, statistics)) as Result<(PartitionedFile, Statistics)> + } else { + Ok((part_file, Statistics::new_unknown(&self.file_schema))) + as Result<(PartitionedFile, Statistics)> + } + }) + .boxed() + .buffered(ctx.config_options().execution.meta_fetch_concurrency); + + let (files, statistics) = + get_statistics_with_limit(files, self.schema(), limit, self.options.collect_stat) + .await?; + + Ok(( + split_files(files, self.options.target_partitions), + statistics, + )) + } + + /// Collects statistics for a given partitioned file. + /// + /// This method first checks if the statistics for the given file are already cached. + /// If they are, it returns the cached statistics. + /// If they are not, it infers the statistics from the file and stores them in the cache. + async fn do_collect_statistics<'a>( + &'a self, + ctx: &SessionState, + store: &Arc, + part_file: &PartitionedFile, + ) -> Result { + let statistics_cache = self.collected_statistics.clone(); + return match statistics_cache + .get_with_extra(&part_file.object_meta.location, &part_file.object_meta) + { + Some(statistics) => Ok(statistics.as_ref().clone()), + None => { + let statistics = self + .options + .format + .infer_stats(ctx, store, self.file_schema.clone(), &part_file.object_meta) + .await?; + statistics_cache.put_with_extra( + &part_file.object_meta.location, + statistics.clone().into(), + &part_file.object_meta, + ); + Ok(statistics) + } + }; + } +} + +#[async_trait] +impl TableProvider for NewListingTable { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.table_schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &SessionState, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result> { + let (mut partitioned_file_lists, statistics) = + self.list_files_for_scan(state, filters, limit).await?; + + let projected_schema = project_schema(&self.schema(), projection)?; + + // if no files need to be read, return an `EmptyExec` + if partitioned_file_lists.is_empty() { + return Ok(Arc::new(EmptyExec::new(projected_schema))); + } + + let output_ordering = self.try_create_output_ordering()?; + match state + .config_options() + .execution + .split_file_groups_by_statistics + .then(|| { + output_ordering.first().map(|output_ordering| { + FileScanConfig::split_groups_by_statistics( + &self.table_schema, + &partitioned_file_lists, + output_ordering, + ) + }) + }) + .flatten() + { + Some(Err(e)) => log::debug!("failed to split file groups by statistics: {e}"), + Some(Ok(new_groups)) => { + if new_groups.len() <= self.options.target_partitions { + partitioned_file_lists = new_groups; + } else { + log::debug!( + "attempted to split file groups by statistics, but there were more file groups than target_partitions; falling back to unordered" + ) + } + } + None => {} // no ordering required + }; + + // extract types of partition columns + let table_partition_cols = self + .options + .table_partition_cols + .iter() + .map(|col| Ok(self.table_schema.field_with_name(&col.0)?.clone())) + .collect::>>()?; + + let filters = if let Some(expr) = conjunction(filters.to_vec()) { + // NOTE: Use the table schema (NOT file schema) here because `expr` may contain + // references to partition columns. + let table_df_schema = self.table_schema.as_ref().clone().to_dfschema()?; + let filters = create_physical_expr(&expr, &table_df_schema, state.execution_props())?; + Some(filters) + } else { + None + }; + + let object_store_url = if let Some(url) = self.table_paths.first() { + url.object_store() + } else { + return Ok(Arc::new(EmptyExec::new(Arc::new(Schema::empty())))); + }; + + // create the execution plan + self.options + .format + .create_physical_plan( + state, + FileScanConfig::new(object_store_url, Arc::clone(&self.file_schema)) + .with_file_groups(partitioned_file_lists) + .with_statistics(statistics) + .with_projection(projection.cloned()) + .with_limit(limit) + .with_output_ordering(output_ordering) + .with_table_partition_cols(table_partition_cols), + filters.as_ref(), + ) + .await + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + let support: Vec<_> = filters + .iter() + .map(|filter| { + if expr_applicable_for_cols( + &self + .options + .table_partition_cols + .iter() + .map(|x| x.0.clone()) + .collect::>(), + filter, + ) { + // if filter can be handled by partition pruning, it is exact + TableProviderFilterPushDown::Exact + } else { + // otherwise, we still might be able to handle the filter with file + // level mechanisms such as Parquet row group pruning. + TableProviderFilterPushDown::Inexact + } + }) + .collect(); + Ok(support) + } +} + +fn create_ordering(schema: &Schema, sort_order: &[Vec]) -> Result> { + let mut all_sort_orders = vec![]; + + for exprs in sort_order { + // Construct PhysicalSortExpr objects from Expr objects: + let mut sort_exprs = vec![]; + for expr in exprs { + match expr { + Expr::Sort(sort) => match sort.expr.as_ref() { + Expr::Column(col) => match expressions::col(&col.name, schema) { + Ok(expr) => { + sort_exprs.push(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !sort.asc, + nulls_first: sort.nulls_first, + }, + }); + } + // Cannot find expression in the projected_schema, stop iterating + // since rest of the orderings are violated + Err(_) => break, + }, + expr => { + return plan_err!( + "Expected single column references in output_ordering, got {expr}" + ); + } + }, + expr => return plan_err!("Expected Expr::Sort in output_ordering, but got {expr}"), + } + } + if !sort_exprs.is_empty() { + all_sort_orders.push(sort_exprs); + } + } + Ok(all_sort_orders) +} diff --git a/src/service/search/datafusion/table_provider/test_util.rs b/src/service/search/datafusion/table_provider/test_util.rs new file mode 100644 index 0000000000..ecee4dc3ba --- /dev/null +++ b/src/service/search/datafusion/table_provider/test_util.rs @@ -0,0 +1,44 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Object store implementation used for testing + +use std::sync::Arc; + +use datafusion::execution::{ + config::SessionConfig, context::SessionState, runtime_env::RuntimeEnv, +}; +use futures::FutureExt; +use object_store::{memory::InMemory, path::Path, ObjectStore}; + +/// Create a test object store with the provided files +pub fn make_test_store_and_state(files: &[(&str, u64)]) -> (Arc, SessionState) { + let memory = InMemory::new(); + + for (name, size) in files { + memory + .put(&Path::from(*name), vec![0; *size as usize].into()) + .now_or_never() + .unwrap() + .unwrap(); + } + + ( + Arc::new(memory), + SessionState::new_with_config_rt(SessionConfig::default(), Arc::new(RuntimeEnv::default())), + ) +} diff --git a/src/service/search/datafusion/arr_descending_udf.rs b/src/service/search/datafusion/udf/arr_descending_udf.rs similarity index 100% rename from src/service/search/datafusion/arr_descending_udf.rs rename to src/service/search/datafusion/udf/arr_descending_udf.rs diff --git a/src/service/search/datafusion/arrcount_udf.rs b/src/service/search/datafusion/udf/arrcount_udf.rs similarity index 100% rename from src/service/search/datafusion/arrcount_udf.rs rename to src/service/search/datafusion/udf/arrcount_udf.rs diff --git a/src/service/search/datafusion/arrindex_udf.rs b/src/service/search/datafusion/udf/arrindex_udf.rs similarity index 100% rename from src/service/search/datafusion/arrindex_udf.rs rename to src/service/search/datafusion/udf/arrindex_udf.rs diff --git a/src/service/search/datafusion/arrjoin_udf.rs b/src/service/search/datafusion/udf/arrjoin_udf.rs similarity index 98% rename from src/service/search/datafusion/arrjoin_udf.rs rename to src/service/search/datafusion/udf/arrjoin_udf.rs index 0287dc5cc3..968b6b487e 100644 --- a/src/service/search/datafusion/arrjoin_udf.rs +++ b/src/service/search/datafusion/udf/arrjoin_udf.rs @@ -28,8 +28,6 @@ use datafusion::{ use datafusion_expr::ColumnarValue; use once_cell::sync::Lazy; -use crate::service::search::datafusion::stringify_json_value; - /// The name of the arrjoin UDF given to DataFusion. pub const ARR_JOIN_UDF_NAME: &str = "arrjoin"; @@ -75,7 +73,7 @@ pub fn arr_join_impl(args: &[ColumnarValue]) -> datafusion::error::Result datafusion::error::Result datafusion::error::Result. + +use config::utils::json; + +use crate::common::meta::functions::ZoFunction; + +pub(crate) mod arr_descending_udf; +pub(crate) mod arrcount_udf; +pub(crate) mod arrindex_udf; +pub(crate) mod arrjoin_udf; +pub(crate) mod arrsort_udf; +pub(crate) mod arrzip_udf; +pub(crate) mod cast_to_arr_udf; +pub(crate) mod date_format_udf; +pub(crate) mod match_udf; +pub(crate) mod regexp_udf; +pub(crate) mod spath_udf; +pub(crate) mod string_to_array_v2_udf; +pub(crate) mod time_range_udf; +pub(crate) mod to_arr_string_udf; +pub(crate) mod transform_udf; + +/// The name of the match UDF given to DataFusion. +pub(crate) const MATCH_UDF_NAME: &str = "str_match"; +/// The name of the match_ignore_case UDF given to DataFusion. +pub(crate) const MATCH_UDF_IGNORE_CASE_NAME: &str = "str_match_ignore_case"; +/// The name of the regex_match UDF given to DataFusion. +pub(crate) const REGEX_MATCH_UDF_NAME: &str = "re_match"; +/// The name of the not_regex_match UDF given to DataFusion. +pub(crate) const REGEX_NOT_MATCH_UDF_NAME: &str = "re_not_match"; + +pub(crate) const DEFAULT_FUNCTIONS: [ZoFunction; 7] = [ + ZoFunction { + name: "match_all_raw", + text: "match_all_raw('v')", + }, + ZoFunction { + name: "match_all_raw_ignore_case", + text: "match_all_raw_ignore_case('v')", + }, + ZoFunction { + name: "match_all", + text: "match_all('v')", + }, + ZoFunction { + name: MATCH_UDF_NAME, + text: "match_all('v')", + }, + ZoFunction { + name: MATCH_UDF_IGNORE_CASE_NAME, + text: "match_all_ignore_case('v')", + }, + ZoFunction { + name: REGEX_MATCH_UDF_NAME, + text: "re_match(field, 'pattern')", + }, + ZoFunction { + name: REGEX_NOT_MATCH_UDF_NAME, + text: "re_not_match(field, 'pattern')", + }, +]; + +pub fn stringify_json_value(field: &json::Value) -> String { + match field { + serde_json::Value::Bool(b) => b.to_string(), + serde_json::Value::Number(n) => match n.as_f64() { + Some(f) => f.to_string(), + None => n.as_i64().unwrap().to_string(), + }, + serde_json::Value::String(s) => s.clone(), + _ => json::to_string(field).expect("failed to stringify json field"), + } +} diff --git a/src/service/search/datafusion/regexp_udf.rs b/src/service/search/datafusion/udf/regexp_udf.rs similarity index 100% rename from src/service/search/datafusion/regexp_udf.rs rename to src/service/search/datafusion/udf/regexp_udf.rs diff --git a/src/service/search/datafusion/spath_udf.rs b/src/service/search/datafusion/udf/spath_udf.rs similarity index 97% rename from src/service/search/datafusion/spath_udf.rs rename to src/service/search/datafusion/udf/spath_udf.rs index cfb01e9dd4..9b092afa8b 100644 --- a/src/service/search/datafusion/spath_udf.rs +++ b/src/service/search/datafusion/udf/spath_udf.rs @@ -28,8 +28,6 @@ use datafusion::{ use datafusion_expr::ColumnarValue; use once_cell::sync::Lazy; -use crate::service::search::datafusion::stringify_json_value; - /// The name of the arrjoin UDF given to DataFusion. pub const SPATH_UDF_NAME: &str = "spath"; @@ -88,7 +86,7 @@ pub fn spath_impl(args: &[ColumnarValue]) -> datafusion::error::Result j #[cfg(test)] mod tests { - use std::sync::Arc; - use datafusion::{ - arrow::{ - array::{Int64Array, StringArray}, - datatypes::{DataType, Field, Schema}, - record_batch::RecordBatch, - }, + arrow::{array::Int64Array, datatypes::Schema, record_batch::RecordBatch}, datasource::MemTable, prelude::SessionContext, }; - use crate::service::search::datafusion::transform_udf::get_udf_vrl; + use super::*; #[tokio::test] async fn vrl_udf_test() {