use std::collections::HashMap;
use std::sync::{Arc, LazyLock};
use itertools::Itertools;
use strum::Display;
use url::Url;
use crate::actions::deletion_vector::DeletionVectorDescriptor;
use crate::error::add_scalar_path_context;
use crate::expressions::{ColumnName, ExpressionRef, PredicateRef, Scalar, StructData};
use crate::schema::{DataType, SchemaRef, StructField, StructType, ToSchema};
use crate::utils::CollectInto;
use crate::{DeltaResult, Error, FileMeta};
#[derive(Debug, Clone, Display)]
#[strum(serialize_all = "snake_case")]
pub enum Operator {
ScanParquet(ScanParquet),
ScanJson(ScanJson),
Values(Values),
Project(Project),
Filter(Filter),
DynamicScan(DynamicScan),
Aggregate(Aggregate),
SemiJoin(SemiJoin),
UnionAll(UnionAll),
}
macro_rules! impl_from_payload_for_operator {
($($variant:ident),+ $(,)?) => {
$(impl From<$variant> for Operator {
fn from(payload: $variant) -> Self {
Operator::$variant(payload)
}
})+
};
}
impl_from_payload_for_operator!(
ScanParquet,
ScanJson,
Values,
Project,
Filter,
DynamicScan,
Aggregate,
SemiJoin,
UnionAll,
);
#[derive(Debug, Clone, PartialEq)]
pub struct ScanFile {
pub meta: FileMeta,
pub file_constants: Vec<Scalar>,
}
impl ScanFile {
pub fn new(meta: FileMeta) -> Self {
Self {
meta,
file_constants: Vec::new(),
}
}
}
impl From<FileMeta> for ScanFile {
fn from(meta: FileMeta) -> Self {
Self::new(meta)
}
}
#[derive(Debug, Clone)]
pub struct ScanParquet {
pub files: Vec<ScanFile>,
pub file_constant_columns: Vec<String>,
pub schema: SchemaRef,
}
#[derive(Debug, Clone)]
pub struct ScanJson {
pub files: Vec<ScanFile>,
pub file_constant_columns: Vec<String>,
pub schema: SchemaRef,
}
#[derive(Debug, Clone)]
pub struct Values {
pub schema: SchemaRef,
pub rows: Vec<Vec<Scalar>>,
}
impl Values {
pub fn new(schema: impl Into<SchemaRef>, rows: Vec<Vec<Scalar>>) -> Self {
Self {
schema: schema.into(),
rows,
}
}
}
impl<T: Into<StructData> + ToSchema> FromIterator<T> for Values {
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
let rows = iter.into_iter().map(|row| row.into().into_parts().1);
Self::new(Arc::new(T::to_schema()), rows.collect())
}
}
impl<T> TryFrom<Values> for Vec<T>
where
T: TryFrom<StructData, Error = Error> + ToSchema,
{
type Error = Error;
fn try_from(Values { schema, rows }: Values) -> DeltaResult<Self> {
rows.into_iter()
.enumerate()
.map(|(index, row)| {
let schema = schema.as_ref().clone();
T::try_from(StructData::from_values_unchecked(schema, row))
.map_err(|error| add_scalar_path_context(error, format!("[{index}]")))
})
.try_collect()
}
}
#[derive(Debug, Clone)]
pub struct Project {
pub expr: ExpressionRef,
pub schema: SchemaRef,
}
#[derive(Debug, Clone)]
pub struct Filter {
pub predicate: PredicateRef,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FileType {
Parquet,
Json,
}
#[derive(Debug, Clone)]
pub struct DynamicScan {
pub schema: SchemaRef,
pub file_type: FileType,
pub base_url: Url,
pub file_constant_columns: Vec<String>,
pub path_column: ColumnName,
pub file_size_column: ColumnName,
pub last_modified_column: ColumnName,
pub dv_column: ColumnName,
}
impl DynamicScan {
#[allow(clippy::too_many_arguments)]
pub fn try_new(
input_schema: &SchemaRef,
output_schema: impl Into<SchemaRef>,
file_type: FileType,
base_url: Url,
file_constant_columns: impl IntoIterator<Item = impl Into<String>>,
path_column: ColumnName,
file_size_column: ColumnName,
last_modified_column: ColumnName,
dv_column: ColumnName,
) -> DeltaResult<Self> {
let schema = output_schema.into();
let file_constant_columns = file_constant_columns
.into_iter()
.map(Into::into)
.collect::<Vec<_>>();
let dynamic_scan = Self {
schema,
file_type,
base_url,
file_constant_columns,
path_column,
file_size_column,
last_modified_column,
dv_column,
};
dynamic_scan.validate_input(input_schema)?;
Ok(dynamic_scan)
}
pub fn validate_input(&self, input_schema: &SchemaRef) -> DeltaResult<()> {
static DELETION_VECTOR_DATA_TYPE: LazyLock<DataType> =
LazyLock::new(|| DataType::from(DeletionVectorDescriptor::to_schema()));
if self.base_url.cannot_be_a_base() || !self.base_url.path().ends_with('/') {
return Err(Error::generic(format!(
"dynamic scan: base URL `{}` must be hierarchical and end in `/`",
self.base_url
)));
}
Self::validate_required_column(input_schema, &self.path_column, &DataType::STRING)?;
Self::validate_required_column(input_schema, &self.file_size_column, &DataType::LONG)?;
Self::validate_required_column(input_schema, &self.last_modified_column, &DataType::LONG)?;
Self::validate_file_constant_columns(
input_schema,
&self.schema,
&self.file_constant_columns,
)?;
let fields = input_schema
.fields_of_path(&self.dv_column)
.map_err(|err| {
Error::generic(format!(
"dynamic scan: deletion-vector column `{}` is invalid: {err}",
self.dv_column
))
})?;
let Some((field, _ancestors)) = fields.split_last() else {
return Err(Error::internal_error("fields_of_path returned no fields"));
};
let expected = &*DELETION_VECTOR_DATA_TYPE;
if field.data_type() != expected {
return Err(Error::generic(format!(
"dynamic scan: deletion-vector column `{}` must have type {expected}, found {}",
self.dv_column,
field.data_type()
)));
}
if !field.is_nullable() {
return Err(Error::generic(format!(
"dynamic scan: deletion-vector column `{}` must be nullable",
self.dv_column
)));
}
Ok(())
}
fn validate_required_column(
schema: &SchemaRef,
column: &ColumnName,
expected_type: &DataType,
) -> DeltaResult<()> {
let fields = schema.fields_of_path(column)?;
let Some((field, ancestors)) = fields.split_last() else {
return Err(Error::internal_error("fields_of_path returned no fields"));
};
if field.data_type() != expected_type {
return Err(Error::generic(format!(
"dynamic scan: column `{column}` must have type {expected_type}, found {}",
field.data_type()
)));
}
if field.is_nullable() || ancestors.iter().any(|field| field.is_nullable()) {
return Err(Error::generic(format!(
"dynamic scan: required column `{column}` is nullable"
)));
}
Ok(())
}
fn validate_file_constant_columns(
input_schema: &SchemaRef,
output_schema: &SchemaRef,
file_constant_columns: &[String],
) -> DeltaResult<()> {
for name in file_constant_columns {
let Some(input_field) = input_schema.field(name) else {
return Err(Error::generic(format!(
"dynamic scan file_constant source: column `{name}` not found; schema has \
{:?}",
Vec::from_iter(input_schema.fields().map(|field| field.name())),
)));
};
if input_field.is_metadata_column() {
return Err(Error::generic(format!(
"dynamic scan file_constant source: column `{name}` is a metadata column"
)));
}
let Some(output_field) = output_schema.field(name) else {
return Err(Error::generic(format!(
"dynamic scan file_constant: column `{name}` not found; schema has {:?}",
Vec::from_iter(output_schema.fields().map(|field| field.name())),
)));
};
if output_field.is_metadata_column() {
return Err(Error::generic(format!(
"dynamic scan file_constant: column `{name}` is a metadata column"
)));
}
if input_field.data_type() != output_field.data_type()
|| input_field.is_nullable() != output_field.is_nullable()
{
return Err(Error::generic(format!(
"dynamic scan file_constant: column `{name}` must have the same type and \
nullability in input and output"
)));
}
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct Aggregate {
pub group_by: Vec<ColumnName>,
pub aggs: Vec<Agg>,
pub schema: SchemaRef,
}
impl Aggregate {
pub fn ungrouped(input_schema: SchemaRef) -> AggregateBuilder {
Self::group_by(input_schema, std::iter::empty::<ColumnName>())
}
pub fn group_by(
input_schema: SchemaRef,
grouping_keys: impl CollectInto<Vec<ColumnName>>,
) -> AggregateBuilder {
AggregateBuilder {
input_schema,
group_by: grouping_keys.collect_into(),
aggs: Vec::new(),
}
}
}
#[derive(Debug, Clone)]
pub enum Agg {
Min(ColumnName),
Max(ColumnName),
Sum(ColumnName),
Count(ColumnName),
CountStar,
MinNonNullBy(NonNullByOperands),
MaxNonNullBy(NonNullByOperands),
}
#[derive(Debug, Clone)]
pub struct NonNullByOperands {
pub value: ColumnName,
pub null_sentinel: ColumnName,
pub key: ColumnName,
}
impl Agg {
pub fn min(value: impl Into<ColumnName>) -> Self {
Self::Min(value.into())
}
pub fn max(value: impl Into<ColumnName>) -> Self {
Self::Max(value.into())
}
pub fn sum(value: impl Into<ColumnName>) -> Self {
Self::Sum(value.into())
}
pub fn count(value: impl Into<ColumnName>) -> Self {
Self::Count(value.into())
}
pub fn count_star() -> Self {
Self::CountStar
}
pub fn min_non_null_by(
value: impl Into<ColumnName>,
null_sentinel: impl Into<ColumnName>,
key: impl Into<ColumnName>,
) -> Self {
Self::MinNonNullBy(NonNullByOperands {
value: value.into(),
null_sentinel: null_sentinel.into(),
key: key.into(),
})
}
pub fn max_non_null_by(
value: impl Into<ColumnName>,
null_sentinel: impl Into<ColumnName>,
key: impl Into<ColumnName>,
) -> Self {
Self::MaxNonNullBy(NonNullByOperands {
value: value.into(),
null_sentinel: null_sentinel.into(),
key: key.into(),
})
}
fn output_field(
&self,
input_schema: &StructType,
alias: Option<String>,
) -> DeltaResult<StructField> {
let resolve = |value: &ColumnName, output_data_type: Option<DataType>, nullable: bool| {
let field = input_schema.field_at(value)?;
let (data_type, metadata) = match output_data_type {
Some(data_type) => (data_type, HashMap::new()),
None => (field.data_type.clone(), field.metadata.clone()),
};
Ok(StructField {
name: alias.clone().unwrap_or_else(|| field.name.clone()),
data_type,
metadata,
nullable,
})
};
match self {
Agg::Min(value) | Agg::Max(value) => resolve(value, None, true),
Agg::Sum(value) => resolve(value, Some(DataType::LONG), true),
Agg::Count(value) => resolve(value, Some(DataType::LONG), false),
Agg::CountStar => Ok(StructField::not_null(
alias.unwrap_or_else(|| "count".to_string()),
DataType::LONG,
)),
Agg::MinNonNullBy(operands) | Agg::MaxNonNullBy(operands) => {
let _ = input_schema.field_at(&operands.key)?;
let _ = input_schema.field_at(&operands.null_sentinel)?;
resolve(&operands.value, None, true)
}
}
}
}
#[derive(Debug)]
pub struct AggregateBuilder {
input_schema: SchemaRef,
group_by: Vec<ColumnName>,
aggs: Vec<(Agg, Option<String>)>,
}
impl AggregateBuilder {
pub fn aggregate(mut self, agg: Agg) -> Self {
self.aggs.push((agg, None));
self
}
pub fn aggregate_as(mut self, agg: Agg, name: impl Into<String>) -> Self {
self.aggs.push((agg, Some(name.into())));
self
}
pub fn min(self, value: impl Into<ColumnName>) -> Self {
self.aggregate(Agg::min(value))
}
pub fn max(self, value: impl Into<ColumnName>) -> Self {
self.aggregate(Agg::max(value))
}
pub fn sum(self, value: impl Into<ColumnName>) -> Self {
self.aggregate(Agg::sum(value))
}
pub fn count(self, value: impl Into<ColumnName>) -> Self {
self.aggregate(Agg::count(value))
}
pub fn count_star(self) -> Self {
self.aggregate(Agg::count_star())
}
pub fn min_non_null_by(
self,
value: impl Into<ColumnName>,
null_sentinel: impl Into<ColumnName>,
key: impl Into<ColumnName>,
) -> Self {
self.aggregate(Agg::min_non_null_by(value, null_sentinel, key))
}
pub fn max_non_null_by(
self,
value: impl Into<ColumnName>,
null_sentinel: impl Into<ColumnName>,
key: impl Into<ColumnName>,
) -> Self {
self.aggregate(Agg::max_non_null_by(value, null_sentinel, key))
}
pub fn build(self) -> DeltaResult<Aggregate> {
let mut fields = Vec::with_capacity(self.group_by.len() + self.aggs.len());
for key in &self.group_by {
fields.push(self.input_schema.field_at(key)?.clone());
}
let mut aggs = Vec::with_capacity(self.aggs.len());
for (agg, alias) in self.aggs {
fields.push(agg.output_field(&self.input_schema, alias)?);
aggs.push(agg);
}
Ok(Aggregate {
group_by: self.group_by,
aggs,
schema: Arc::new(StructType::try_new(fields)?),
})
}
}
impl TryFrom<AggregateBuilder> for Aggregate {
type Error = Error;
fn try_from(builder: AggregateBuilder) -> DeltaResult<Self> {
builder.build()
}
}
#[derive(Debug, Clone)]
pub struct SemiJoin {
pub inverted: bool,
pub probe_keys: Vec<ColumnName>,
pub build_keys: Vec<ColumnName>,
}
#[derive(Debug, Clone)]
pub struct UnionAll;
#[cfg(test)]
mod tests {
use delta_kernel_derive::{IntoStructData, ToSchema, TryFromStructData};
use super::*;
use crate::expressions::column_name;
use crate::schema::{DataType, MetadataValue, StructField};
use crate::unit_test_utils::assert_result_error_with_message;
fn schema(fields: &[(&str, bool)]) -> SchemaRef {
Arc::new(StructType::new_unchecked(fields.iter().map(
|(name, nullable)| StructField::new(*name, DataType::LONG, *nullable),
)))
}
#[test]
fn output_lists_group_keys_then_aggregates_in_order() {
let input = schema(&[("g", false), ("a", true), ("b", true)]);
let agg = Aggregate::group_by(input, [column_name!("g")])
.max(column_name!("a"))
.min(column_name!("b"))
.build()
.unwrap();
let names: Vec<&str> = agg.schema.fields().map(|f| f.name().as_str()).collect();
assert_eq!(names, ["g", "a", "b"]);
}
#[test]
fn output_fields_preserve_input_field_metadata() {
let metadata = [("k", MetadataValue::Number(7))];
let input = Arc::new(StructType::new_unchecked([
StructField::not_null("g", DataType::LONG).with_metadata(metadata.clone()),
StructField::not_null("a", DataType::LONG).with_metadata(metadata.clone()),
StructField::not_null("s", DataType::LONG).with_metadata(metadata),
]));
let agg = Aggregate::group_by(input, [column_name!("g")])
.max(column_name!("a"))
.sum(column_name!("s"))
.build()
.unwrap();
let key = agg.schema.field("g").unwrap();
assert!(!key.nullable);
assert_eq!(key.metadata()["k"], MetadataValue::Number(7));
let max = agg.schema.field("a").unwrap();
assert!(max.nullable);
assert_eq!(max.metadata()["k"], MetadataValue::Number(7));
assert!(agg.schema.field("s").unwrap().metadata().is_empty());
}
#[rstest::rstest]
#[case::min(Agg::min(column_name!("a")), "a", true)]
#[case::max(Agg::max(column_name!("a")), "a", true)]
#[case::sum(Agg::sum(column_name!("a")), "a", true)]
#[case::count(Agg::count(column_name!("a")), "a", false)]
#[case::count_star(Agg::count_star(), "count", false)]
#[case::min_non_null_by(
Agg::min_non_null_by(column_name!("a"), column_name!("s"), column_name!("v")),
"a",
true
)]
#[case::max_non_null_by(
Agg::max_non_null_by(column_name!("a"), column_name!("s"), column_name!("v")),
"a",
true
)]
fn agg_output_nullability(
#[case] agg: Agg,
#[case] name: &str,
#[case] nullable: bool,
#[values(true, false)] value_nullable: bool,
) {
let input = schema(&[("a", value_nullable), ("s", true), ("v", true)]);
let built = Aggregate::ungrouped(input).aggregate(agg).build().unwrap();
let field = built.schema.field(name).unwrap();
assert_eq!(field.nullable, nullable);
assert_eq!(field.data_type(), &DataType::LONG);
}
#[test]
fn alias_overrides_default_output_name() {
let input = schema(&[("a", true)]);
let agg = Aggregate::group_by(input, [])
.aggregate_as(Agg::max(column_name!("a")), "a_max")
.build()
.unwrap();
assert!(agg.schema.field("a_max").is_some());
assert!(agg.schema.field("a").is_none());
}
#[test]
fn duplicate_output_names_are_rejected() {
let input = schema(&[("a", true)]);
let result = Aggregate::group_by(input, [])
.min(column_name!("a"))
.max(column_name!("a"))
.build();
assert_result_error_with_message(result, "Duplicate field name");
}
#[test]
fn distinct_aliases_resolve_min_max_collision() {
let input = schema(&[("a", true)]);
let agg = Aggregate::group_by(input, [])
.aggregate_as(Agg::min(column_name!("a")), "a_min")
.aggregate_as(Agg::max(column_name!("a")), "a_max")
.build()
.unwrap();
let names: Vec<&str> = agg.schema.fields().map(|f| f.name().as_str()).collect();
assert_eq!(names, ["a_min", "a_max"]);
}
#[rstest::rstest]
#[case::missing_value_column(false)]
#[case::missing_group_key(true)]
fn build_rejects_missing_column(#[case] missing_in_key: bool) {
let input = schema(&[("a", true)]);
let (keys, value) = if missing_in_key {
(vec![column_name!("missing")], column_name!("a"))
} else {
(vec![], column_name!("missing"))
};
let result = Aggregate::group_by(input, keys).max(value).build();
assert_result_error_with_message(result, "missing");
}
#[test]
fn build_rejects_missing_non_null_by_key() {
let input = schema(&[("a", true)]);
let result = Aggregate::group_by(input, [])
.max_non_null_by(
column_name!("a"),
column_name!("a"),
column_name!("missing"),
)
.build();
assert_result_error_with_message(result, "missing");
}
#[test]
fn build_rejects_missing_non_null_by_sentinel_column() {
let input = schema(&[("a", true), ("v", true)]);
let result = Aggregate::group_by(input, [])
.max_non_null_by(
column_name!("a"),
column_name!("missing"),
column_name!("v"),
)
.build();
assert_result_error_with_message(result, "missing");
}
#[derive(Clone, Debug, PartialEq, ToSchema, IntoStructData, TryFromStructData)]
struct Address {
city: String,
}
#[derive(Clone, Debug, PartialEq, ToSchema, IntoStructData, TryFromStructData)]
struct Person {
id: i32,
address: Address,
}
#[test]
fn values_from_iter_peels_top_level_and_keeps_nested_struct() {
let values = Values::from_iter([Person {
id: 1,
address: Address { city: "NYC".into() },
}]);
assert_eq!(
values
.schema
.fields()
.map(|f| f.name().as_str())
.collect::<Vec<_>>(),
["id", "address"]
);
assert_eq!(values.rows.len(), 1);
assert_eq!(values.rows[0].len(), 2);
assert_eq!(values.rows[0][0], Scalar::Integer(1));
let Scalar::Struct(address) = &values.rows[0][1] else {
panic!("expected nested Struct for address");
};
assert_eq!(address.values(), &[Scalar::String("NYC".into())]);
}
#[test]
fn values_from_iter_empty_still_carries_schema() {
let values: Values = std::iter::empty::<Person>().collect();
assert!(values.rows.is_empty());
assert_eq!(values.schema.num_fields(), 2);
}
#[test]
fn values_round_trips_through_vec() {
let people = vec![
Person {
id: 1,
address: Address { city: "NYC".into() },
},
Person {
id: 2,
address: Address { city: "SF".into() },
},
];
let values = Values::from_iter(people.clone());
assert_eq!(Vec::<Person>::try_from(values).unwrap(), people);
}
#[test]
fn values_conversion_adds_row_index_to_error_path() {
let mut values = Values::from_iter([Person {
id: 1,
address: Address { city: "NYC".into() },
}]);
values.rows[0][0] = Scalar::from("not an integer");
assert_result_error_with_message(
Vec::<Person>::try_from(values),
"[0].id: expected i32, found string",
);
}
}