From 863ac080966a68561bf2b5777a00960982a97db6 Mon Sep 17 00:00:00 2001 From: Andrew Gazelka Date: Wed, 20 Nov 2024 01:24:41 -0800 Subject: [PATCH] [REFACTOR] connect: `to_daft_*` use ref instead of value --- src/daft-connect/src/translation/expr.rs | 12 ++++----- .../translation/expr/unresolved_function.rs | 9 +++---- src/daft-connect/src/translation/literal.rs | 26 +++++++++---------- .../src/translation/logical_plan/aggregate.rs | 4 +-- .../src/translation/logical_plan/project.rs | 2 +- 5 files changed, 26 insertions(+), 27 deletions(-) diff --git a/src/daft-connect/src/translation/expr.rs b/src/daft-connect/src/translation/expr.rs index 299a770d3c..bcbadf9737 100644 --- a/src/daft-connect/src/translation/expr.rs +++ b/src/daft-connect/src/translation/expr.rs @@ -9,12 +9,12 @@ use crate::translation::to_daft_literal; mod unresolved_function; -pub fn to_daft_expr(expression: Expression) -> eyre::Result { - if let Some(common) = expression.common { +pub fn to_daft_expr(expression: &Expression) -> eyre::Result { + if let Some(common) = &expression.common { warn!("Ignoring common metadata for relation: {common:?}; not yet implemented"); }; - let Some(expr) = expression.expr_type else { + let Some(expr) = &expression.expr_type else { bail!("Expression is required"); }; @@ -35,7 +35,7 @@ pub fn to_daft_expr(expression: Expression) -> eyre::Result { warn!("Ignoring is_metadata_column {is_metadata_column} for attribute expressions; not yet implemented"); } - Ok(daft_dsl::col(unparsed_identifier)) + Ok(daft_dsl::col(unparsed_identifier.as_str())) } spark_expr::ExprType::UnresolvedFunction(f) => { unresolved_to_daft_expr(f).wrap_err("Failed to handle unresolved function") @@ -49,7 +49,7 @@ pub fn to_daft_expr(expression: Expression) -> eyre::Result { expr, name, metadata, - } = *alias; + } = &**alias; let Some(expr) = expr else { bail!("Alias expr is required"); @@ -63,7 +63,7 @@ pub fn to_daft_expr(expression: Expression) -> eyre::Result { bail!("Alias metadata is not yet supported; got {metadata:?}"); } - let child = to_daft_expr(*expr)?; + let child = to_daft_expr(expr)?; let name = Arc::from(name.as_str()); diff --git a/src/daft-connect/src/translation/expr/unresolved_function.rs b/src/daft-connect/src/translation/expr/unresolved_function.rs index fc8e7f39b1..ffb8c802ce 100644 --- a/src/daft-connect/src/translation/expr/unresolved_function.rs +++ b/src/daft-connect/src/translation/expr/unresolved_function.rs @@ -1,11 +1,10 @@ use daft_core::count_mode::CountMode; -use daft_schema::dtype::DataType; use eyre::{bail, Context}; use spark_connect::expression::UnresolvedFunction; use crate::translation::to_daft_expr; -pub fn unresolved_to_daft_expr(f: UnresolvedFunction) -> eyre::Result { +pub fn unresolved_to_daft_expr(f: &UnresolvedFunction) -> eyre::Result { let UnresolvedFunction { function_name, arguments, @@ -13,13 +12,13 @@ pub fn unresolved_to_daft_expr(f: UnresolvedFunction) -> eyre::Result = arguments.into_iter().map(to_daft_expr).try_collect()?; + let arguments: Vec<_> = arguments.iter().map(to_daft_expr).try_collect()?; - if is_distinct { + if *is_distinct { bail!("Distinct not yet supported"); } - if is_user_defined_function { + if *is_user_defined_function { bail!("User-defined functions not yet supported"); } diff --git a/src/daft-connect/src/translation/literal.rs b/src/daft-connect/src/translation/literal.rs index 0eff840702..5f45632e4b 100644 --- a/src/daft-connect/src/translation/literal.rs +++ b/src/daft-connect/src/translation/literal.rs @@ -2,18 +2,18 @@ use eyre::{bail, WrapErr}; use spark_connect::expression::{literal::LiteralType, Literal}; // todo(test): add tests for this esp in Python -pub fn to_daft_literal(literal: Literal) -> eyre::Result { - let Some(literal) = literal.literal_type else { +pub fn to_daft_literal(literal: &Literal) -> eyre::Result { + let Some(literal) = &literal.literal_type else { bail!("Literal is required"); }; match literal { LiteralType::Array(_) => bail!("Array literals not yet supported"), LiteralType::Binary(bytes) => Ok(daft_dsl::lit(bytes.as_slice())), - LiteralType::Boolean(b) => Ok(daft_dsl::lit(b)), + LiteralType::Boolean(b) => Ok(daft_dsl::lit(*b)), LiteralType::Byte(b) => { // todo(correctness): is this signed or unsigned? - let b = i8::try_from(b).wrap_err_with(|| format!("Byte value {b} is out of range"))?; + let b = i8::try_from(*b).wrap_err_with(|| format!("Byte value {b} is out of range"))?; let b = i32::from(b); Ok(daft_dsl::lit::(b)) @@ -21,18 +21,18 @@ pub fn to_daft_literal(literal: Literal) -> eyre::Result { LiteralType::CalendarInterval(_) => { bail!("Calendar interval literals not yet supported") } - LiteralType::Date(d) => Ok(daft_dsl::lit(d)), + LiteralType::Date(d) => Ok(daft_dsl::lit(*d)), LiteralType::DayTimeInterval(_) => { bail!("Day-time interval literals not yet supported") } LiteralType::Decimal(_) => bail!("Decimal literals not yet supported"), - LiteralType::Double(d) => Ok(daft_dsl::lit(d)), + LiteralType::Double(d) => Ok(daft_dsl::lit(*d)), LiteralType::Float(f) => { - let f = f64::from(f); + let f = f64::from(*f); Ok(daft_dsl::lit(f)) } - LiteralType::Integer(i) => Ok(daft_dsl::lit(i)), - LiteralType::Long(l) => Ok(daft_dsl::lit(l)), + LiteralType::Integer(i) => Ok(daft_dsl::lit(*i)), + LiteralType::Long(l) => Ok(daft_dsl::lit(*l)), LiteralType::Map(_) => bail!("Map literals not yet supported"), LiteralType::Null(_) => { // todo(correctness): is it ok to assume type is i32 here? @@ -40,19 +40,19 @@ pub fn to_daft_literal(literal: Literal) -> eyre::Result { } LiteralType::Short(s) => { let short = - i16::try_from(s).wrap_err_with(|| format!("Short value {s} is out of range"))?; + i16::try_from(*s).wrap_err_with(|| format!("Short value {s} is out of range"))?; Ok(daft_dsl::lit(i32::from(short))) } - LiteralType::String(s) => Ok(daft_dsl::lit(s)), + LiteralType::String(s) => Ok(daft_dsl::lit(s.as_str())), LiteralType::Struct(_) => bail!("Struct literals not yet supported"), LiteralType::Timestamp(ts) => { // todo(correctness): is it ok that the type is different logically? - Ok(daft_dsl::lit(ts)) + Ok(daft_dsl::lit(*ts)) } LiteralType::TimestampNtz(ts) => { // todo(correctness): is it ok that the type is different logically? - Ok(daft_dsl::lit(ts)) + Ok(daft_dsl::lit(*ts)) } LiteralType::YearMonthInterval(_) => { bail!("Year-month interval literals not yet supported") diff --git a/src/daft-connect/src/translation/logical_plan/aggregate.rs b/src/daft-connect/src/translation/logical_plan/aggregate.rs index 42a003f0df..cbb2aa6348 100644 --- a/src/daft-connect/src/translation/logical_plan/aggregate.rs +++ b/src/daft-connect/src/translation/logical_plan/aggregate.rs @@ -24,12 +24,12 @@ pub fn aggregate(aggregate: spark_connect::Aggregate) -> eyre::Result = grouping_expressions - .into_iter() + .iter() .map(to_daft_expr) .try_collect()?; let aggregate_expressions: Vec<_> = aggregate_expressions - .into_iter() + .iter() .map(to_daft_expr) .try_collect()?; diff --git a/src/daft-connect/src/translation/logical_plan/project.rs b/src/daft-connect/src/translation/logical_plan/project.rs index 1a5614cbc2..3096b7f313 100644 --- a/src/daft-connect/src/translation/logical_plan/project.rs +++ b/src/daft-connect/src/translation/logical_plan/project.rs @@ -18,7 +18,7 @@ pub fn project(project: Project) -> eyre::Result { let plan = to_logical_plan(*input)?; - let daft_exprs: Vec<_> = expressions.into_iter().map(to_daft_expr).try_collect()?; + let daft_exprs: Vec<_> = expressions.iter().map(to_daft_expr).try_collect()?; let plan = plan.select(daft_exprs)?;