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

feat: add custom table provider (#3846)

Browse Source
This commit is contained in:
Hengfei Yang
2024-06-29 20:08:55 +08:00
committed by GitHub
parent 2518f6aa6c
commit ac4b804494
23 changed files with 1434 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/handler/http/request/status/mod.rs
+1 -1
View File
@@ -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},
},
};
src/service/promql/engine.rs
+3 -2
View File
@@ -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())]),
)?
src/service/search/datafusion/exec.rs
+25 -23
View File
@@ -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());
}
src/service/search/datafusion/mod.rs
+2 -71
View File
@@ -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')",
},
];
src/service/search/datafusion/table_provider/helpers.rs
+887
View File
@@ -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<PartitionedFile>,
n: usize,
) -> Vec<Vec<PartitionedFile>> {
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<Vec<ObjectMeta>>,
}
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<Path>)> {
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<Path>,
) -> Result<Vec<Partition>> {
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<Partition>,
filters: &[Expr],
partition_cols: &[(String, DataType)],
) -> Result<Vec<Partition>> {
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::<Result<_, _>>()?;
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<ArrayRef> {
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<Path> {
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<BoxStream<'a, Result<PartitionedFile>>> {
// 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::<Result<Vec<_>>>()?;
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<Vec<&'a str>>
where
I: IntoIterator<Item = &'a str>,
{
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<impl Iterator<Item = &'b str> + '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::<Vec<_>>()
.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::<Vec<_>>()
.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::<Vec<_>>()
.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")),
);
}
}
src/service/search/datafusion/table_provider/mod.rs
+379
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<ListingTableUrl>,
/// 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<Self> {
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<FileStatisticsCache>) -> 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<Vec<LexOrdering>> {
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<usize>,
) -> Result<(Vec<Vec<PartitionedFile>>, 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<dyn ObjectStore>,
part_file: &PartitionedFile,
) -> Result<Statistics> {
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<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
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::<Result<Vec<_>>>()?;
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<Vec<TableProviderFilterPushDown>> {
let support: Vec<_> = filters
.iter()
.map(|filter| {
if expr_applicable_for_cols(
&self
.options
.table_partition_cols
.iter()
.map(|x| x.0.clone())
.collect::<Vec<_>>(),
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<Expr>]) -> Result<Vec<LexOrdering>> {
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)
}
src/service/search/datafusion/table_provider/test_util.rs
+44
View File
@@ -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<InMemory>, 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())),
)
}
src/service/search/datafusion/arr_descending_udf.rs → src/service/search/datafusion/udf/arr_descending_udf.rs
View File
src/service/search/datafusion/arrcount_udf.rs → src/service/search/datafusion/udf/arrcount_udf.rs
View File
src/service/search/datafusion/arrindex_udf.rs → src/service/search/datafusion/udf/arrindex_udf.rs
View File
src/service/search/datafusion/arrjoin_udf.rs → src/service/search/datafusion/udf/arrjoin_udf.rs
+1 -3
View File
@@ -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<Column
let mut join_arrs = vec![];
if let json::Value::Array(field1) = arr_field1 {
field1.iter().for_each(|field| {
let field = stringify_json_value(field);
let field = super::stringify_json_value(field);
join_arrs.push(field);
});
}
src/service/search/datafusion/arrsort_udf.rs → src/service/search/datafusion/udf/arrsort_udf.rs
View File
src/service/search/datafusion/arrzip_udf.rs → src/service/search/datafusion/udf/arrzip_udf.rs
+2 -4
View File
@@ -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 arrzip UDF given to DataFusion.
pub const ARR_ZIP_UDF_NAME: &str = "arrzip";
@@ -83,8 +81,8 @@ pub fn arr_zip_impl(args: &[ColumnarValue]) -> datafusion::error::Result<Columna
{
// Field1 and field2 can be of different types
zip(field1.iter(), field2.iter()).for_each(|(field1, field2)| {
let field1 = stringify_json_value(field1);
let field2 = stringify_json_value(field2);
let field1 = super::stringify_json_value(field1);
let field2 = super::stringify_json_value(field2);
zipped_arrs.push(format!("{field1}{delim}{field2}"));
});
src/service/search/datafusion/cast_to_arr_udf.rs → src/service/search/datafusion/udf/cast_to_arr_udf.rs
+1 -3
View File
@@ -29,8 +29,6 @@ use datafusion::{
use datafusion_expr::ColumnarValue;
use once_cell::sync::Lazy;
use crate::service::search::datafusion::stringify_json_value;
/// The name of the cast_to_arr UDF given to DataFusion.
pub const CAST_TO_ARR_UDF_NAME: &str = "cast_to_arr";
@@ -74,7 +72,7 @@ pub fn cast_to_arr_impl(args: &[ColumnarValue]) -> datafusion::error::Result<Col
json::from_str(string).expect("Failed to deserialize the field into an array");
if let json::Value::Array(arr) = arr {
arr.iter().for_each(|field| {
let field = stringify_json_value(field);
let field = super::stringify_json_value(field);
if !field.is_empty() {
list_builder.values().append_value(field);
}
src/service/search/datafusion/date_format_udf.rs → src/service/search/datafusion/udf/date_format_udf.rs
View File
src/service/search/datafusion/match_udf.rs → src/service/search/datafusion/udf/match_udf.rs
View File
src/service/search/datafusion/udf/mod.rs
+86
View File
@@ -0,0 +1,86 @@
// 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 <http://www.gnu.org/licenses/>.
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"),
}
}
src/service/search/datafusion/regexp_udf.rs → src/service/search/datafusion/udf/regexp_udf.rs
View File
src/service/search/datafusion/spath_udf.rs → src/service/search/datafusion/udf/spath_udf.rs
+1 -3
View File
@@ -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<ColumnarV
}
}
if found {
let field = stringify_json_value(&field);
let field = super::stringify_json_value(&field);
Some(field)
} else {
None
src/service/search/datafusion/string_to_array_v2_udf.rs → src/service/search/datafusion/udf/string_to_array_v2_udf.rs
View File
src/service/search/datafusion/time_range_udf.rs → src/service/search/datafusion/udf/time_range_udf.rs
View File
src/service/search/datafusion/to_arr_string_udf.rs → src/service/search/datafusion/udf/to_arr_string_udf.rs
View File
src/service/search/datafusion/transform_udf.rs → src/service/search/datafusion/udf/transform_udf.rs
+2 -8
View File
@@ -254,19 +254,13 @@ pub fn apply_vrl_fn(runtime: &mut Runtime, program: vrl::compiler::Program) -> 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() {