From 83e3e0bcc69fe3a9fed03ec8002d4d81d7bc56cc Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Sun, 25 Oct 2026 02:57:36 +0200 Subject: [PATCH 1/7] Implement misc scalar functions: uuid, rand, uniform, version, typeof and the rest Registers and binds the misc functions used by the v4.2.0 goldens misc-functions.sql, random.sql and current_database_catalog.sql: version, current_database/schema/catalog, input_file_name and the block functions, monotonically_increasing_id, spark_partition_id, typeof, rand (random), randn, uuid, uniform and randstr. Deterministic ones bind as constants (version from the kernels' Spark line; database/catalog are the engine's single namespace) or as an analyzer literal rewrite (typeof, resolved from the argument's analyzed type like Spark's RuntimeReplaceable). The nondeterministic ones own one value stream per partition: (partition << 33) + row for monotonically_increasing_id, the partition index for spark_partition_id, and a seeded XORShiftRandom stream (seed + partition, advanced with the partition's rows) for rand/randn/uniform/randstr/uuid, so seeded results match Spark exactly. The partition index reaches the evaluator from the driver's task context through the project/filter operators; plans stay pure data and the plan codec is unchanged. Unseeded calls draw from a per-query seed mixed with a per-occurrence salt (Spark's ResolveRandomSeed assigns each its own seed). Spark-compatibility notes: uuid draws its 128 bits from XORShift rather than Spark's MersenneTwister (shape and per-seed determinism match, not the values); version reports the release line without Spark's git revision; input_file_name and the block functions return Spark's no-file values (the evaluator has no file context yet); shuffle stays known-unsupported (it needs an array type the engine cannot represent). Statements needing array/map/struct/interval types (typeof of those, shuffle) and FROM VALUES are listed as skipped: the engine has neither yet. The two multi-row goldens depend on Spark CI's partitioning (two slices, then sorted output); the engine runs tiny scans in one partition, so those rows are asserted sorted where the value multiset agrees, and per-partition at the operator level otherwise. Anything unverified: randstr/uuid/shuffle seeded values against Spark beyond the goldens' first rows; exact error text of INVALID_PARAMETER for uuid/shuffle non-literal seeds beyond the probed cases. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/context.rs | 19 + crates/krishiv-exec/src/driver.rs | 14 +- crates/krishiv-exec/src/operators/expr.rs | 20 + crates/krishiv-exec/src/operators/filter.rs | 16 +- crates/krishiv-exec/src/operators/mod.rs | 8 + crates/krishiv-exec/src/operators/project.rs | 16 +- crates/krishiv-exec/src/operators/tests.rs | 83 +++ crates/krishiv-exec/src/planner/eval/arith.rs | 18 +- crates/krishiv-exec/src/planner/eval/cast.rs | 4 +- crates/krishiv-exec/src/planner/eval/cond.rs | 40 +- crates/krishiv-exec/src/planner/eval/func.rs | 11 +- crates/krishiv-exec/src/planner/eval/logic.rs | 23 +- crates/krishiv-exec/src/planner/eval/mod.rs | 49 +- .../krishiv-exec/src/planner/eval/strings.rs | 22 +- crates/krishiv-exec/src/planner/subquery.rs | 2 +- crates/krishiv-exec/src/runner.rs | 5 +- crates/krishiv-functions/src/registry/defs.rs | 1 + crates/krishiv-functions/src/registry/misc.rs | 336 +++++++++++ crates/krishiv-functions/src/registry/mod.rs | 1 + .../krishiv-functions/src/registry/tests.rs | 278 ++++++++- crates/krishiv-functions/src/scalar/misc.rs | 160 +++++ crates/krishiv-functions/src/scalar/mod.rs | 95 ++- crates/krishiv-functions/src/scalar/random.rs | 533 +++++++++++++++++ .../src/scalar/tests/misc.rs | 138 +++++ .../krishiv-functions/src/scalar/tests/mod.rs | 3 + .../src/scalar/tests/random.rs | 375 ++++++++++++ crates/krishiv-plan/src/analyzer/exprs.rs | 28 + crates/krishiv-plan/src/tests/analyzer.rs | 5 +- .../krishiv-session/tests/misc_functions.rs | 481 +++++++++++++++ .../analysis/registry-matrix-manifest.json | 556 +++++++++++++++++- 30 files changed, 3251 insertions(+), 89 deletions(-) create mode 100644 crates/krishiv-functions/src/registry/misc.rs create mode 100644 crates/krishiv-functions/src/scalar/misc.rs create mode 100644 crates/krishiv-functions/src/scalar/random.rs create mode 100644 crates/krishiv-functions/src/scalar/tests/misc.rs create mode 100644 crates/krishiv-functions/src/scalar/tests/random.rs create mode 100644 crates/krishiv-session/tests/misc_functions.rs diff --git a/crates/krishiv-exec/src/context.rs b/crates/krishiv-exec/src/context.rs index e9b58aa0f..3c511db03 100644 --- a/crates/krishiv-exec/src/context.rs +++ b/crates/krishiv-exec/src/context.rs @@ -398,6 +398,10 @@ pub struct TaskContext { name: &'static str, memory: OnceLock, dependency: Option, + /// The driver's index in its pipeline run (Spark's partition index for + /// the partition the driver computes): stateful functions own one value + /// stream per partition. One driver always reads partition 0. + partition: u32, } impl fmt::Debug for TaskContext { @@ -426,9 +430,24 @@ impl TaskContext { name, memory: OnceLock::new(), dependency: None, + partition: 0, } } + /// The same context for the driver `partition` (its index in the + /// pipeline run). + #[must_use] + pub(crate) fn with_partition(mut self, partition: u32) -> Self { + self.partition = partition; + self + } + + /// The driver's index in its pipeline run (set when the driver is + /// spawned; 0 for a single driver). + pub fn partition(&self) -> u32 { + self.partition + } + pub(crate) fn with_dependency(mut self, state: GlobalState) -> Self { self.dependency = Some(state); self diff --git a/crates/krishiv-exec/src/driver.rs b/crates/krishiv-exec/src/driver.rs index 356a349ef..5f2308e76 100644 --- a/crates/krishiv-exec/src/driver.rs +++ b/crates/krishiv-exec/src/driver.rs @@ -119,14 +119,15 @@ fn since(executor: &dyn Executor, start: Instant) -> Duration { executor.now().saturating_duration_since(start) } -/// Spawns one driver of `run`. -pub(crate) fn spawn_driver(run: Arc) -> TaskHandle { +/// Spawns one driver of `run`: `partition` is the driver's index in the +/// run (see [`TaskContext::with_partition`]). +pub(crate) fn spawn_driver(run: Arc, partition: u32) -> TaskHandle { let pipeline = run.pipeline.index; let query = Arc::clone(&run.query); let executor = Arc::clone(&run.executor); let q = Arc::clone(&query); let body = async move { - if let Err(e) = drive(run).await { + if let Err(e) = drive(run, partition).await { q.fail(e); } }; @@ -145,8 +146,8 @@ pub(crate) fn spawn_driver(run: Arc) -> TaskHandle { executor.spawn(query.class(), Box::pin(timed)) } -async fn drive(run: Arc) -> Result<()> { - Driver::new(run)?.run().await +async fn drive(run: Arc, partition: u32) -> Result<()> { + Driver::new(run, partition)?.run().await } /// Times every poll of a driver: the stall detector. A poll is a @@ -287,7 +288,7 @@ struct Driver { } impl Driver { - fn new(run: Arc) -> Result { + fn new(run: Arc, partition: u32) -> Result { let owner = SpillOwner::new(); let ctx = |id, name| { TaskContext::new( @@ -297,6 +298,7 @@ impl Driver { id, name, ) + .with_partition(partition) }; let p = &run.pipeline; let (source_id, source) = &p.source; diff --git a/crates/krishiv-exec/src/operators/expr.rs b/crates/krishiv-exec/src/operators/expr.rs index a9aea6646..fbd1666ab 100644 --- a/crates/krishiv-exec/src/operators/expr.rs +++ b/crates/krishiv-exec/src/operators/expr.rs @@ -20,6 +20,16 @@ pub trait Predicate: Send + Sync + fmt::Debug { /// Any error fails the query (for example a division by zero under /// ANSI mode). fn evaluate(&self, batch: &RecordBatch) -> Result; + /// As [`Self::evaluate`], in the executing partition `partition` + /// (Spark's partition index). The default runs [`Self::evaluate`], for + /// the stateless predicates. + /// + /// # Errors + /// As [`Self::evaluate`]. + fn evaluate_in(&self, batch: &RecordBatch, partition: u32) -> Result { + let _ = partition; + self.evaluate(batch) + } } /// A scalar expression evaluated over a batch, for [`Project`](super::Project). @@ -33,6 +43,16 @@ pub trait ScalarExpr: Send + Sync + fmt::Debug { /// # Errors /// Any error fails the query. fn evaluate(&self, batch: &RecordBatch) -> Result; + /// As [`Self::evaluate`], in the executing partition `partition` + /// (Spark's partition index). The default runs [`Self::evaluate`], for + /// the stateless expressions. + /// + /// # Errors + /// As [`Self::evaluate`]. + fn evaluate_in(&self, batch: &RecordBatch, partition: u32) -> Result { + let _ = partition; + self.evaluate(batch) + } } /// A [`Predicate`] from a closure. diff --git a/crates/krishiv-exec/src/operators/filter.rs b/crates/krishiv-exec/src/operators/filter.rs index 5c850bda1..c65ac360d 100644 --- a/crates/krishiv-exec/src/operators/filter.rs +++ b/crates/krishiv-exec/src/operators/filter.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use arrow_schema::SchemaRef; use krishiv_types::batch::Batch; +use super::PartitionState; use super::expr::Predicate; -use super::no_state; use crate::context::TaskContext; use crate::error::Result; use crate::operator::{LocalState, Operator, OperatorOutput}; @@ -38,15 +38,21 @@ impl Operator for Filter { Arc::clone(&self.schema) } - fn local_state(&self, _ctx: &TaskContext) -> Result { - Ok(no_state()) + fn local_state(&self, ctx: &TaskContext) -> Result { + Ok(Box::new(PartitionState { + partition: ctx.partition(), + })) } - fn execute(&self, input: &Batch, _local: &mut LocalState) -> Result { + fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { if input.num_rows() == 0 { return Ok(OperatorOutput::NeedMoreInput(None)); } - let mask = self.predicate.evaluate(input.data())?; + let partition = local + .downcast_ref::() + .map(|s| s.partition) + .unwrap_or(0); + let mask = self.predicate.evaluate_in(input.data(), partition)?; let out = input.filter(&mask)?; Ok(OperatorOutput::NeedMoreInput( (out.num_rows() > 0).then_some(out), diff --git a/crates/krishiv-exec/src/operators/mod.rs b/crates/krishiv-exec/src/operators/mod.rs index 4ebc36e92..ddb7c1f6b 100644 --- a/crates/krishiv-exec/src/operators/mod.rs +++ b/crates/krishiv-exec/src/operators/mod.rs @@ -84,5 +84,13 @@ pub(crate) fn no_state() -> LocalState { Box::new(()) } +/// Per-driver local state of an operator evaluating row expressions: the +/// driver's partition index (see [`TaskContext::partition`](crate::context::TaskContext::partition)). +#[derive(Debug, Clone, Copy)] +pub(crate) struct PartitionState { + /// The driver's partition index. + pub(crate) partition: u32, +} + #[cfg(all(test, not(loom)))] mod tests; diff --git a/crates/krishiv-exec/src/operators/project.rs b/crates/krishiv-exec/src/operators/project.rs index a6dbdd3ba..d29758a3e 100644 --- a/crates/krishiv-exec/src/operators/project.rs +++ b/crates/krishiv-exec/src/operators/project.rs @@ -8,8 +8,8 @@ use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions}; use arrow_schema::{Field, Schema, SchemaRef}; use krishiv_types::batch::Batch; +use super::PartitionState; use super::expr::ScalarExpr; -use super::no_state; use crate::context::TaskContext; use crate::error::{Result, internal}; use crate::operator::{LocalState, Operator, OperatorOutput}; @@ -138,15 +138,21 @@ impl Operator for Project { Arc::clone(&self.schema) } - fn local_state(&self, _ctx: &TaskContext) -> Result { - Ok(no_state()) + fn local_state(&self, ctx: &TaskContext) -> Result { + Ok(Box::new(PartitionState { + partition: ctx.partition(), + })) } - fn execute(&self, input: &Batch, _local: &mut LocalState) -> Result { + fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { let rows = input.num_rows(); if rows == 0 { return Ok(OperatorOutput::NeedMoreInput(None)); } + let partition = local + .downcast_ref::() + .map(|s| s.partition) + .unwrap_or(0); let data = input.data(); let columns = self .items @@ -158,7 +164,7 @@ impl Operator for Project { .get(*index) .map(Arc::clone) .ok_or_else(|| internal(format!("project index {index} out of range"))), - ProjectItem::Expr { expr, .. } => expr.evaluate(data), + ProjectItem::Expr { expr, .. } => expr.evaluate_in(data, partition), } }) .collect::>>()?; diff --git a/crates/krishiv-exec/src/operators/tests.rs b/crates/krishiv-exec/src/operators/tests.rs index d628fb75a..744e53df5 100644 --- a/crates/krishiv-exec/src/operators/tests.rs +++ b/crates/krishiv-exec/src/operators/tests.rs @@ -965,3 +965,86 @@ fn expand_rejects_bad_plans() { .unwrap_or_else(|| unreachable!("an error")); assert_eq!(err.error_class(), ErrorClass::InternalError); } + +/// A stateful function sees the driver's partition: `rand(0)` draws +/// partition 0's stream on one driver and partition 1's on another, and +/// the same partition draws the same stream on a fresh evaluation. +#[test] +fn project_threads_the_drivers_partition_to_stateful_functions() { + use arrow_array::NullArray; + use arrow_array::types::Float64Type; + use krishiv_common::config::Config; + use krishiv_plan::{ + Expr, ExprNode, FunctionId, FunctionKind, Literal, ResolvedExpr, ResolvedFunction, + }; + + use crate::planner::eval::{Compiled, EvalEnv, Input, compile}; + + // `rand(0)` analyzed: a non-nullable double of a folded seed. + let call = |seed: i32| { + let f = ResolvedFunction { + id: FunctionId::Registry { + name: "rand".to_owned(), + id: 0, + }, + kind: FunctionKind::Scalar, + args: vec![Arc::new(Expr::new(ExprNode::Literal(Literal::Integer( + seed, + ))))], + distinct: false, + filter: None, + data_type: krishiv_types::DataType::Double, + nullable: false, + deterministic: false, + }; + Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f))) + }; + let one_row = || { + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Null, true)])); + let data = RecordBatch::try_new(schema, vec![Arc::new(NullArray::new(1))]) + .unwrap_or_else(|e| unreachable!("{e}")); + Batch::new(data) + }; + let project = |seed: i32| { + let config = Config::new(); + let env = EvalEnv::new(&config, Arc::from("test")); + let schema = Arc::new(Schema::empty()); + let compiled = compile( + &call(seed), + Input { + layout: &[], + schema: &schema, + env: &env, + subqueries: &[], + }, + ) + .unwrap_or_else(|e| unreachable!("{e}")); + let input = one_row(); + Project::try_new( + &input.schema(), + vec![ProjectItem::expr( + "r", + Arc::new(Compiled(compiled)) as Arc, + )], + ) + .unwrap_or_else(|e| unreachable!("{e}")) + }; + let run = |seed: i32, partition: u32| { + let p = project(seed); + let (_root, ctx) = ctx(); + let ctx = ctx.with_partition(partition); + let mut local = p.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); + let out = out_batch( + p.execute(&one_row(), &mut local) + .unwrap_or_else(|e| unreachable!("{e}")), + ) + .unwrap_or_else(|| unreachable!("a row")); + out.data().column(0).as_primitive::().value(0) + }; + // Seed 0, partition 0 and 1: Spark's first draws. + assert_eq!(run(0, 0), 0.7604953758285915); + assert_eq!(run(0, 1), 0.6363787615254752); + // A fresh evaluation draws the same stream again: the stream is a pure + // function of (seed, partition, row). + assert_eq!(run(0, 0), 0.7604953758285915); +} diff --git a/crates/krishiv-exec/src/planner/eval/arith.rs b/crates/krishiv-exec/src/planner/eval/arith.rs index 07d1f37ee..0864dd26f 100644 --- a/crates/krishiv-exec/src/planner/eval/arith.rs +++ b/crates/krishiv-exec/src/planner/eval/arith.rs @@ -151,8 +151,11 @@ impl PhysExpr for Binary { self.nullable } - fn eval(&self, batch: &RecordBatch) -> Result { - let (mut l, r) = (self.l.eval(batch)?, self.r.eval(batch)?); + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let (mut l, r) = ( + self.l.eval(batch, partition)?, + self.r.eval(batch, partition)?, + ); let scalar = l.is_scalar() && r.is_scalar(); if scalar { l = Value::Array(Arc::clone(l.array())); @@ -213,8 +216,8 @@ impl PhysExpr for Negative { self.e.nullable() } - fn eval(&self, batch: &RecordBatch) -> Result { - let v = self.e.eval(batch)?; + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let v = self.e.eval(batch, partition)?; let a = v.array(); let ctx = &self.kernels; macro_rules! int { @@ -294,8 +297,11 @@ impl PhysExpr for IntegralDivide { true } - fn eval(&self, batch: &RecordBatch) -> Result { - let (mut l, r) = (self.l.eval(batch)?, self.r.eval(batch)?); + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let (mut l, r) = ( + self.l.eval(batch, partition)?, + self.r.eval(batch, partition)?, + ); let scalar = l.is_scalar() && r.is_scalar(); if scalar { l = Value::Array(Arc::clone(l.array())); diff --git a/crates/krishiv-exec/src/planner/eval/cast.rs b/crates/krishiv-exec/src/planner/eval/cast.rs index 4763b528c..819675105 100644 --- a/crates/krishiv-exec/src/planner/eval/cast.rs +++ b/crates/krishiv-exec/src/planner/eval/cast.rs @@ -123,8 +123,8 @@ impl PhysExpr for Cast { self.nullable } - fn eval(&self, batch: &RecordBatch) -> Result { - let v = self.e.eval(batch)?; + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let v = self.e.eval(batch, partition)?; let out = apply(&self.parts, v.array())?; Ok(match v { Value::Scalar(_) => Value::Scalar(out), diff --git a/crates/krishiv-exec/src/planner/eval/cond.rs b/crates/krishiv-exec/src/planner/eval/cond.rs index 943c5f6be..b4d750804 100644 --- a/crates/krishiv-exec/src/planner/eval/cond.rs +++ b/crates/krishiv-exec/src/planner/eval/cond.rs @@ -93,7 +93,7 @@ impl PhysExpr for When { || self.branches.iter().any(|(_, v)| v.nullable()) } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { let rows = batch.num_rows(); // `rest`: the rows no branch took yet, as a batch and their row // numbers in `batch`. @@ -106,7 +106,7 @@ impl PhysExpr for When { if rest_rows.is_empty() { break; } - let c = cond.eval(&rest)?.into_array(rest.num_rows())?; + let c = cond.eval(&rest, partition)?.into_array(rest.num_rows())?; let c = c .as_boolean_opt() .ok_or_else(|| internal("CASE condition"))?; @@ -121,7 +121,9 @@ impl PhysExpr for When { let mask = BooleanArray::new(taken.clone(), None); let chosen = filter_record_batch(&rest, &mask).map_err(|e| internal(&e.to_string()))?; let v = cast_null( - value.eval(&chosen)?.into_array(chosen.num_rows())?, + value + .eval(&chosen, partition)? + .into_array(chosen.num_rows())?, &self.out, ); let part = parts.len(); @@ -143,7 +145,10 @@ impl PhysExpr for When { } if !rest_rows.is_empty() { let v = match &self.otherwise { - Some(e) => cast_null(e.eval(&rest)?.into_array(rest.num_rows())?, &self.out), + Some(e) => cast_null( + e.eval(&rest, partition)?.into_array(rest.num_rows())?, + &self.out, + ), None => new_null_array(&self.out, rest_rows.len()), }; let part = parts.len(); @@ -262,12 +267,12 @@ impl PhysExpr for NullIf { true } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { let rows = batch.num_rows(); // `value` is evaluated once: the comparison casts its array, the // result keeps its values. - let a = self.value.eval(batch)?; - let b = self.other.eval(batch)?; + let a = self.value.eval(batch, partition)?; + let b = self.other.eval(batch, partition)?; let scalar = a.is_scalar() && b.is_scalar(); let n = if scalar { 1 } else { rows }; let a_full = a.into_array(n)?; @@ -348,10 +353,10 @@ impl PhysExpr for Nanvl { self.a.nullable() || self.b.nullable() } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { let rows = batch.num_rows(); // `a` is evaluated once; NULL is not NaN. - let a = self.a.eval(batch)?.into_array(rows)?; + let a = self.a.eval(batch, partition)?.into_array(rows)?; let nan = match a.data_type() { ArrowType::Float32 => BooleanBuffer::from_iter( a.as_primitive::() @@ -369,7 +374,7 @@ impl PhysExpr for Nanvl { return Ok(Value::Array(a)); } // `b` sees only the NaN rows, so `nanvl(1, 1 / 0)` never raises. - let b = super::eval_where(self.b.as_ref(), batch, &nan)?; + let b = super::eval_where(self.b.as_ref(), batch, &nan, partition)?; let out = zip(&BooleanArray::new(nan, None), &b, &a).map_err(|e| internal(&e.to_string()))?; Ok(Value::Array(out)) @@ -457,7 +462,10 @@ mod tests { /// Evaluates `e` over `batch` as a column of `batch.num_rows()` rows. fn values(e: &Arc, batch: &RecordBatch) -> ArrayRef { - e.eval(batch).unwrap().into_array(batch.num_rows()).unwrap() + e.eval(batch, 0) + .unwrap() + .into_array(batch.num_rows()) + .unwrap() } /// `1 / 0`: `NULL` without ANSI, `DIVIDE_BY_ZERO` with it. @@ -529,7 +537,7 @@ mod tests { ); let e = coalesce(&f, vec![one_over_zero(&env), lit(&Literal::Double(1.0))]).unwrap(); if ansi { - let err = e.eval(&batch).unwrap_err(); + let err = e.eval(&batch, 0).unwrap_err(); assert_eq!(err.error_class(), ErrorClass::DivideByZero); } else { assert_eq!( @@ -666,7 +674,7 @@ mod tests { .unwrap(); match want { None => { - let err = e.eval(&batch).unwrap_err(); + let err = e.eval(&batch, 0).unwrap_err(); assert_eq!(err.error_class(), ErrorClass::CastInvalidInput); assert_eq!(err.sql_state(), Some("22018")); } @@ -736,9 +744,9 @@ mod tests { self.inner.nullable() } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { self.n.fetch_add(1, Ordering::SeqCst); - self.inner.eval(batch) + self.inner.eval(batch, partition) } } @@ -868,7 +876,7 @@ mod tests { ) .unwrap(); if ansi { - let err = e.eval(&batch).unwrap_err(); + let err = e.eval(&batch, 0).unwrap_err(); assert_eq!(err.error_class(), ErrorClass::DivideByZero); } else { assert!(values(&e, &batch).as_primitive::().is_null(0)); diff --git a/crates/krishiv-exec/src/planner/eval/func.rs b/crates/krishiv-exec/src/planner/eval/func.rs index e469091f7..45c8e2d28 100644 --- a/crates/krishiv-exec/src/planner/eval/func.rs +++ b/crates/krishiv-exec/src/planner/eval/func.rs @@ -51,7 +51,8 @@ pub(super) fn compile( .with_kernel(input.env.kernels) .with_time_zone(Arc::new(zone)) .with_query_start(input.env.query_start_micros) - .with_user(Arc::clone(&input.env.user)); + .with_user(Arc::clone(&input.env.user)) + .with_seed_salt(Arc::clone(&input.env.seed_salt)); let function = bind_scalar(name, &bind_args, &opts).map_err(|err| { if err.error_class() == krishiv_common::error::ErrorClass::KrishivUnsupportedFeature { unsupported_expr(e) @@ -98,18 +99,20 @@ impl PhysExpr for Call { self.nullable } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { let args = self .args .iter() .map(|a| { - Ok(match a.eval(batch)? { + Ok(match a.eval(batch, partition)? { Value::Array(v) => Arg::Array(v), Value::Scalar(v) => Arg::Scalar(v), }) }) .collect::>>()?; - let out = self.function.invoke(&args, batch.num_rows())?; + let out = self + .function + .invoke_in(&args, batch.num_rows(), partition)?; Ok(Value::Array(if self.view { to_view(&out)? } else { out })) } } diff --git a/crates/krishiv-exec/src/planner/eval/logic.rs b/crates/krishiv-exec/src/planner/eval/logic.rs index def638285..ec380350d 100644 --- a/crates/krishiv-exec/src/planner/eval/logic.rs +++ b/crates/krishiv-exec/src/planner/eval/logic.rs @@ -62,8 +62,8 @@ impl PhysExpr for AndOr { self.l.nullable() || self.r.nullable() } - fn eval(&self, batch: &RecordBatch) -> Result { - let lv = self.l.eval(batch)?; + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let lv = self.l.eval(batch, partition)?; if lv.is_scalar() && batch.num_rows() > 0 { // A literal left side decides for every row or for none. let l = lv @@ -87,9 +87,9 @@ impl PhysExpr for AndOr { let decided = if self.and { &l_ok & &!&lb } else { &l_ok & &lb }; let open = !&decided; let rv = if open.count_set_bits() == rows { - self.r.eval(batch)? + self.r.eval(batch, partition)? } else { - Value::Array(super::eval_where(self.r.as_ref(), batch, &open)?) + Value::Array(super::eval_where(self.r.as_ref(), batch, &open, partition)?) }; let scalar = lv.is_scalar() && rv.is_scalar(); let rows = if scalar { 1 } else { rows }; @@ -133,8 +133,8 @@ impl PhysExpr for Not { self.0.nullable() } - fn eval(&self, batch: &RecordBatch) -> Result { - let v = self.0.eval(batch)?; + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let v = self.0.eval(batch, partition)?; let a = v .array() .as_boolean_opt() @@ -167,8 +167,8 @@ impl PhysExpr for IsNull { false } - fn eval(&self, batch: &RecordBatch) -> Result { - let v = self.e.eval(batch)?; + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let v = self.e.eval(batch, partition)?; let a = v.array(); let valid = match (a.data_type(), a.logical_nulls()) { (ArrowType::Null, _) => BooleanBuffer::new_unset(a.len()), @@ -235,8 +235,11 @@ impl PhysExpr for Compare { self.kind.is_some() && (self.l.nullable() || self.r.nullable()) } - fn eval(&self, batch: &RecordBatch) -> Result { - let (mut l, r) = (self.l.eval(batch)?, self.r.eval(batch)?); + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let (mut l, r) = ( + self.l.eval(batch, partition)?, + self.r.eval(batch, partition)?, + ); let scalar = l.is_scalar() && r.is_scalar(); if scalar { // Kernels take at most one literal side. diff --git a/crates/krishiv-exec/src/planner/eval/mod.rs b/crates/krishiv-exec/src/planner/eval/mod.rs index 71fa27969..52de45a49 100644 --- a/crates/krishiv-exec/src/planner/eval/mod.rs +++ b/crates/krishiv-exec/src/planner/eval/mod.rs @@ -31,6 +31,7 @@ mod strings; use std::fmt; use std::sync::Arc; +use std::sync::atomic::AtomicU64; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, UInt32Array, new_null_array}; @@ -90,11 +91,13 @@ pub(crate) fn broadcast(s: &ArrayRef, rows: usize) -> Result { /// `e` over the rows of `batch` selected by `mask`, as a column of every /// row (NULL where not selected): `e` sees only the selected rows, so it -/// raises only for them. +/// raises only for them. `partition` is the executing partition's index +/// (see [`PhysExpr::eval`]). pub(crate) fn eval_where( e: &dyn PhysExpr, batch: &RecordBatch, mask: &arrow_buffer::BooleanBuffer, + partition: u32, ) -> Result { let rows = batch.num_rows(); let selected = mask.count_set_bits(); @@ -104,7 +107,7 @@ pub(crate) fn eval_where( let chosen = arrow_select::filter::filter_record_batch(batch, &BooleanArray::new(mask.clone(), None)) .map_err(|err| internal(&err.to_string()))?; - let values = e.eval(&chosen)?.into_array(selected)?; + let values = e.eval(&chosen, partition)?.into_array(selected)?; scatter(values.as_ref(), mask) } @@ -144,8 +147,13 @@ pub(crate) trait PhysExpr: Send + Sync + fmt::Debug { fn data_type(&self) -> &ArrowType; /// Whether a value may be NULL. fn nullable(&self) -> bool; - /// Its values over `batch`. - fn eval(&self, batch: &RecordBatch) -> Result; + /// Its values over `batch` in the executing partition `partition` + /// (Spark's partition index): a stateful function (`rand`, `uuid`, + /// `monotonically_increasing_id`, ...) owns one value stream per + /// partition, and this evaluates the next values of that stream. The + /// same query runs each partition's stream from its start, so one + /// partition's values never depend on another's rows. + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result; } /// What every expression of one query is evaluated with. @@ -163,6 +171,9 @@ pub(crate) struct EvalEnv { pub(crate) query_start_micros: i64, /// The session user (`current_user()`). pub(crate) user: Arc, + /// Counts the unseeded random binds of one plan instantiation (see + /// [`BindOptions::seed_salt`](krishiv_functions::scalar::BindOptions::seed_salt)). + pub(crate) seed_salt: Arc, } impl EvalEnv { @@ -186,6 +197,7 @@ impl EvalEnv { kernels: KernelCtx::new(), query_start_micros, user, + seed_salt: Arc::new(AtomicU64::new(0)), } } @@ -197,6 +209,7 @@ impl EvalEnv { kernels: KernelCtx::new(), query_start_micros: env.query_start_micros, user: Arc::clone(&env.user), + seed_salt: Arc::new(AtomicU64::new(0)), } } @@ -381,7 +394,7 @@ impl PhysExpr for Column { self.nullable } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, _partition: u32) -> Result { batch .columns() .get(self.index) @@ -414,7 +427,7 @@ impl PhysExpr for Lit { self.value.logical_null_count() > 0 } - fn eval(&self, _batch: &RecordBatch) -> Result { + fn eval(&self, _batch: &RecordBatch, _partition: u32) -> Result { Ok(Value::Scalar(Arc::clone(&self.value))) } } @@ -468,18 +481,22 @@ pub(crate) fn literal_array(l: &Literal) -> Result { pub(crate) struct Compiled(pub(crate) Arc); impl Compiled { - /// The values over `batch`. When evaluation fails, the error is the + /// The values over `batch` in `partition`. When evaluation fails, the error is the /// first failing row's, as Spark's row-at-a-time evaluation raises it: /// a batch evaluates sub-expressions one after another (and each lazy /// `CASE` branch over its own rows), so the first error found need not /// be the first row's. Only the failing path pays for the search. - fn values(&self, batch: &RecordBatch) -> Result { + fn values(&self, batch: &RecordBatch, partition: u32) -> Result { let rows = batch.num_rows(); - match self.0.eval(batch).and_then(|v| v.into_array(rows)) { + match self + .0 + .eval(batch, partition) + .and_then(|v| v.into_array(rows)) + { Ok(v) => Ok(v), Err(e) if rows > 1 => { for i in 0..rows { - self.0.eval(&batch.slice(i, 1))?; + self.0.eval(&batch.slice(i, 1), partition)?; } Err(e) } @@ -490,7 +507,11 @@ impl Compiled { impl Predicate for Compiled { fn evaluate(&self, batch: &RecordBatch) -> Result { - let v = self.values(batch)?; + ::evaluate_in(self, batch, 0) + } + + fn evaluate_in(&self, batch: &RecordBatch, partition: u32) -> Result { + let v = self.values(batch, partition)?; v.as_boolean_opt() .cloned() .ok_or_else(|| internal("a predicate is not boolean")) @@ -507,7 +528,11 @@ impl ScalarExpr for Compiled { } fn evaluate(&self, batch: &RecordBatch) -> Result { - self.values(batch) + ::evaluate_in(self, batch, 0) + } + + fn evaluate_in(&self, batch: &RecordBatch, partition: u32) -> Result { + self.values(batch, partition) } } diff --git a/crates/krishiv-exec/src/planner/eval/strings.rs b/crates/krishiv-exec/src/planner/eval/strings.rs index 49e2b2a82..fe91a1f7c 100644 --- a/crates/krishiv-exec/src/planner/eval/strings.rs +++ b/crates/krishiv-exec/src/planner/eval/strings.rs @@ -80,9 +80,10 @@ struct Like { pub(super) fn like(mut args: Vec>) -> Result> { let escape = if args.len() == 3 { let e = args.pop().ok_or_else(|| internal("LIKE arity"))?; - let v = e.eval(&RecordBatch::new_empty(Arc::new( - arrow_schema::Schema::empty(), - ))); + let v = e.eval( + &RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())), + 0, + ); match v { Ok(Value::Scalar(s)) if is_string(s.data_type()) && s.is_valid(0) => { let text = str_value(s.as_ref(), 0); @@ -114,8 +115,11 @@ impl PhysExpr for Like { self.s.nullable() || self.pattern.nullable() } - fn eval(&self, batch: &RecordBatch) -> Result { - let (s, p) = (self.s.eval(batch)?, self.pattern.eval(batch)?); + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { + let (s, p) = ( + self.s.eval(batch, partition)?, + self.pattern.eval(batch, partition)?, + ); let scalar = s.is_scalar() && p.is_scalar(); let rows = if scalar { 1 } else { batch.num_rows() }; let (sa, pa) = (s.array(), p.array()); @@ -176,11 +180,11 @@ impl PhysExpr for Concat { self.nullable } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, partition: u32) -> Result { let values = self .args .iter() - .map(|a| a.eval(batch)) + .map(|a| a.eval(batch, partition)) .collect::>>()?; let scalar = values.iter().all(Value::is_scalar); let arrays = values @@ -230,7 +234,7 @@ mod tests { true } - fn eval(&self, batch: &RecordBatch) -> Result { + fn eval(&self, batch: &RecordBatch, _partition: u32) -> Result { Ok(Value::Array(Arc::clone(batch.column(self.0)))) } } @@ -282,7 +286,7 @@ mod tests { true, ) .unwrap(); - let Value::Array(out) = e.eval(&batch).unwrap() else { + let Value::Array(out) = e.eval(&batch, 0).unwrap() else { panic!("a column in, a column out"); }; let out: Vec<_> = out.as_string_view().iter().collect(); diff --git a/crates/krishiv-exec/src/planner/subquery.rs b/crates/krishiv-exec/src/planner/subquery.rs index 77323d579..bbc65bb52 100644 --- a/crates/krishiv-exec/src/planner/subquery.rs +++ b/crates/krishiv-exec/src/planner/subquery.rs @@ -141,7 +141,7 @@ impl PhysExpr for SubqueryValue { self.kind == Kind::Scalar } - fn eval(&self, _batch: &RecordBatch) -> Result { + fn eval(&self, _batch: &RecordBatch, _partition: u32) -> Result { let v = self .slot .0 diff --git a/crates/krishiv-exec/src/runner.rs b/crates/krishiv-exec/src/runner.rs index d0ced223f..cdf3ccb35 100644 --- a/crates/krishiv-exec/src/runner.rs +++ b/crates/krishiv-exec/src/runner.rs @@ -432,7 +432,10 @@ fn start_group( query: Arc::clone(query), executor: Arc::clone(executor), }); - handles.extend((0..drivers).map(|_| spawn_driver(Arc::clone(&run)))); + handles.extend((0..drivers).map(|i| { + let partition = u32::try_from(i).unwrap_or(u32::MAX); + spawn_driver(Arc::clone(&run), partition) + })); runs.push(run); } Ok((runs, handles)) diff --git a/crates/krishiv-functions/src/registry/defs.rs b/crates/krishiv-functions/src/registry/defs.rs index de824f8a3..2972609ef 100644 --- a/crates/krishiv-functions/src/registry/defs.rs +++ b/crates/krishiv-functions/src/registry/defs.rs @@ -643,5 +643,6 @@ pub(super) fn specs() -> Vec { out.extend(super::bitwise::specs()); out.extend(super::predicate::specs()); out.extend(super::number::specs()); + out.extend(super::misc::specs()); out } diff --git a/crates/krishiv-functions/src/registry/misc.rs b/crates/krishiv-functions/src/registry/misc.rs new file mode 100644 index 000000000..f294fc8e9 --- /dev/null +++ b/crates/krishiv-functions/src/registry/misc.rs @@ -0,0 +1,336 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The misc scalar signatures: `version`, `current_database` +//! (`current_schema`), `current_catalog`, the `input_file_*` placeholders, +//! `monotonically_increasing_id`, `spark_partition_id`, `typeof` and the +//! seeded random functions (`rand` / `random`, `randn`, `uuid`, `uniform`, +//! `randstr`). +//! +//! Ground truth: Spark 4.2's `misc.scala` (`SparkVersion`, +//! `CurrentDatabase`, `CurrentCatalog`, `InputFileName`, `Uuid`), +//! `randomExpressions.scala` (`Rand`, `Randn`, `Uniform`, `RandStr`, +//! `NondeterministicUnaryRDG`), `MonotonicallyIncreasingID.scala`, +//! `SparkPartitionID.scala`, and the `random.sql`, +//! `current_database_catalog.sql` and `misc-functions.sql` goldens (checked +//! against a real Spark 4.2.0 run where the goldens were ambiguous). +//! +//! Arity, determinism and availability come from Spark's catalog (the +//! `FunctionDef` assembly reads them off the TSV row); only the implicit +//! casts (`Inputs`) and the result types live here. Every rule mirrors an +//! error the goldens pin: `rand`/`randn` and `randstr`/`uniform` seeds must +//! be foldable, `uuid`'s seed a plain int/long literal (Spark's +//! `expressionToSeed`), and `uniform`/`randstr` bounds the types Spark's +//! implicit casts accept. + +use krishiv_common::error::{ + DatatypeMismatchNonFoldableInput, DatatypeMismatchUnexpectedInputType, + InvalidParameterValueLong, KrishivError, SeedExpressionIsUnfoldable, +}; +use krishiv_types::DataType; + +use super::defs::{Spec, scalar}; +use crate::signature::{ArgType, CallTypes, Constant, Inputs}; + +/// `version`, `current_database` (`current_schema`), `current_catalog`, +/// `input_file_name`: a non-nullable string; `input_file_block_start` / +/// `input_file_block_length`: a non-nullable `bigint`; +/// `monotonically_increasing_id`: a non-nullable `bigint`; +/// `spark_partition_id`: a non-nullable `int`. +fn string(c: &CallTypes<'_>) -> Result { + Ok(ArgType::new(DataType::String, c.any_nullable())) +} + +fn long(c: &CallTypes<'_>) -> Result { + Ok(ArgType::new(DataType::Long, c.any_nullable())) +} + +fn int(c: &CallTypes<'_>) -> Result { + Ok(ArgType::new(DataType::Integer, c.any_nullable())) +} + +/// `typeof(x)`: a non-nullable string. The name itself is resolved by the +/// analyzer (a literal rewrite); the rule only types the call. +fn typeof_result(c: &CallTypes<'_>) -> Result { + string(c) +} + +/// Whether `t` is an `int` or a `bigint` (what every non-`uuid` seed takes). +fn is_seed(t: Option<&DataType>) -> bool { + matches!(t, Some(DataType::Integer | DataType::Long) | None) +} + +/// The whole call as Spark quotes it in errors (`rand(id)`). +fn call_sql(name: &str, c: &CallTypes<'_>) -> String { + let args: Vec = (0..c.args().len()) + .map(|i| c.sql(i).unwrap_or_default()) + .collect(); + format!("\"{name}({})\"", args.join(", ")) +} + +/// `rand(seed?)` / `random(seed?)` / `randn(seed?)`: a non-nullable +/// `double`. The seed must be an `int`/`bigint` +/// (`DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE`, as Spark's `Rand` reports +/// it) and foldable (`SEED_EXPRESSION_IS_UNFOLDABLE`); an untyped `NULL` +/// reads as 0. +fn rand_result(name: &'static str, c: &CallTypes<'_>) -> Result { + if !c.args().is_empty() { + if !is_seed(c.data_type(0)) { + return Err(DatatypeMismatchUnexpectedInputType { + sql_expr: call_sql(name, c), + param_index: "first".to_owned(), + required_type: "(\"INT\" or \"BIGINT\")".to_owned(), + input_sql: format!("\"{}\"", c.sql(0).unwrap_or_default()), + input_type: c + .data_type(0) + .map_or_else(|| "\"VOID\"".to_owned(), |t| format!("\"{}\"", t.sql())), + } + .into()); + } + if c.constant(0).is_none() { + return Err(SeedExpressionIsUnfoldable { + seed_expr: format!("\"{}\"", c.sql(0).unwrap_or_default()), + expr_with_seed: call_sql(name, c), + } + .into()); + } + } + Ok(ArgType::new(DataType::Double, false)) +} + +fn rand(c: &CallTypes<'_>) -> Result { + rand_result("rand", c) +} + +fn randn(c: &CallTypes<'_>) -> Result { + rand_result("randn", c) +} + +/// `uuid(seed?)`: a non-nullable string. The seed must be a plain `int` / +/// `bigint` literal (Spark's `expressionToSeed`); a NULL literal draws from +/// the query seed. +fn uuid(c: &CallTypes<'_>) -> Result { + if !c.args().is_empty() { + let literal = c.is_literal(0) + && matches!( + (c.data_type(0), c.constant(0)), + (None, Some(Constant::Null)) | (Some(DataType::Integer | DataType::Long), Some(_)) + ); + if !literal { + return Err(InvalidParameterValueLong { + function_name: "`UUID`".to_owned(), + parameter: "`seed`".to_owned(), + invalid_value: format!("\"{}\"", c.sql(0).unwrap_or_default()), + } + .into()); + } + } + Ok(ArgType::new(DataType::String, false)) +} + +/// Whether `t` is numeric or an untyped `NULL` (what `uniform`'s bounds +/// take; Spark casts the `NULL` to `DOUBLE`). +fn is_uniform_bound(t: Option<&DataType>) -> bool { + t.is_none_or(|t| { + matches!( + t, + DataType::Byte + | DataType::Short + | DataType::Integer + | DataType::Long + | DataType::Float + | DataType::Double + | DataType::Decimal(_) + ) + }) +} + +/// The wider of two integral types (Spark's `UpCastRule` order). +fn wider_integral(left: &DataType, right: &DataType) -> DataType { + let rank = |t: &DataType| match t { + DataType::Byte => 0, + DataType::Short => 1, + DataType::Integer => 2, + _ => 3, + }; + if rank(left) >= rank(right) { + left.clone() + } else { + right.clone() + } +} + +/// Whether `left` is wider than `right` (Spark's +/// `DecimalType.isWiderThan`). +fn is_wider_decimal(left: (u8, u8), right: (u8, u8)) -> bool { + let (p1, s1) = (i32::from(left.0), i32::from(left.1)); + let (p2, s2) = (i32::from(right.0), i32::from(right.1)); + p1 - s1 >= p2 - s2 && s1 >= s2 && p1 >= p2 +} + +/// `uniform(min, max[, seed])`: Spark's `Uniform.dataType` (an untyped +/// `NULL` reads as `DOUBLE`), nullable when an argument may be. A NULL +/// bound is a NULL result; a NULL seed reads as 0. +fn uniform(c: &CallTypes<'_>) -> Result { + for i in 0..3 { + if i >= c.args().len() { + continue; + } + let bound = i < 2; + if bound && !is_uniform_bound(c.data_type(i)) || !bound && !is_seed(c.data_type(i)) { + return Err(DatatypeMismatchUnexpectedInputType { + sql_expr: call_sql("uniform", c), + param_index: ["first", "second", "third"][i].to_owned(), + required_type: if bound { + "\"NUMERIC\"".to_owned() + } else { + "(\"INT\" or \"BIGINT\")".to_owned() + }, + input_sql: format!("\"{}\"", c.sql(i).unwrap_or_default()), + input_type: c + .data_type(i) + .map_or_else(|| "\"VOID\"".to_owned(), |t| format!("\"{}\"", t.sql())), + } + .into()); + } + } + for (i, name) in [(0, "min"), (1, "max"), (2, "seed")] { + if i >= c.args().len() { + continue; + } + match c.constant(i) { + Some(Constant::Error(e)) => return Err(e.clone()), + None => { + return Err(DatatypeMismatchNonFoldableInput { + sql_expr: call_sql("uniform", c), + input_name: format!("`{name}`"), + input_type: "integer or floating-point".to_owned(), + input_expr: format!("\"{}\"", c.sql(i).unwrap_or_default()), + } + .into()); + } + Some(_) => {} + } + } + let num = |t: Option<&DataType>| t.cloned().unwrap_or(DataType::Double); + let (min, max) = (num(c.data_type(0)), num(c.data_type(1))); + let data_type = match (&min, &max) { + (a, b) + if matches!( + a, + DataType::Byte | DataType::Short | DataType::Integer | DataType::Long + ) && matches!( + b, + DataType::Byte | DataType::Short | DataType::Integer | DataType::Long + ) => + { + wider_integral(a, b) + } + (DataType::Double, _) | (_, DataType::Double) => DataType::Double, + (DataType::Float, _) | (_, DataType::Float) => DataType::Float, + (DataType::Decimal(left), DataType::Decimal(right)) + if is_wider_decimal( + (left.precision(), left.scale()), + (right.precision(), right.scale()), + ) => + { + min + } + (_, d @ DataType::Decimal(_)) | (d @ DataType::Decimal(_), _) => d.clone(), + _ => DataType::Double, + }; + Ok(ArgType::new(data_type, c.any_nullable())) +} + +/// Whether `t` is a type Spark implicitly casts to `INT` (the numerics and +/// strings, or an untyped `NULL`): what `randstr`'s length takes. +fn is_length(t: Option<&DataType>) -> bool { + t.is_none_or(|t| { + matches!( + t, + DataType::Byte + | DataType::Short + | DataType::Integer + | DataType::Long + | DataType::Float + | DataType::Double + | DataType::Decimal(_) + | DataType::String + ) + }) +} + +/// `randstr(length[, seed])`: a non-nullable string. The length takes what +/// Spark casts to `INT`; the seed an `int`/`bigint`. Both must be foldable. +fn randstr(c: &CallTypes<'_>) -> Result { + if !is_length(c.data_type(0)) { + return Err(DatatypeMismatchUnexpectedInputType { + sql_expr: call_sql("randstr", c), + param_index: "first".to_owned(), + required_type: "\"INT\"".to_owned(), + input_sql: format!("\"{}\"", c.sql(0).unwrap_or_default()), + input_type: c + .data_type(0) + .map_or_else(|| "\"VOID\"".to_owned(), |t| format!("\"{}\"", t.sql())), + } + .into()); + } + if c.args().len() > 1 && !is_seed(c.data_type(1)) { + return Err(DatatypeMismatchUnexpectedInputType { + sql_expr: call_sql("randstr", c), + param_index: "second".to_owned(), + required_type: "(\"INT\" or \"BIGINT\")".to_owned(), + input_sql: format!("\"{}\"", c.sql(1).unwrap_or_default()), + input_type: c + .data_type(1) + .map_or_else(|| "\"VOID\"".to_owned(), |t| format!("\"{}\"", t.sql())), + } + .into()); + } + for (i, name) in [(0, "length"), (1, "seed")] { + if i >= c.args().len() { + continue; + } + match c.constant(i) { + Some(Constant::Error(e)) => return Err(e.clone()), + None => { + return Err(DatatypeMismatchNonFoldableInput { + sql_expr: call_sql("randstr", c), + input_name: format!("`{name}`"), + input_type: "integer".to_owned(), + input_expr: format!("\"{}\"", c.sql(i).unwrap_or_default()), + } + .into()); + } + Some(_) => {} + } + } + Ok(ArgType::new(DataType::String, false)) +} + +/// Every implemented misc function. +pub(super) fn specs() -> Vec { + let mut out = scalar( + &[ + "version", + "current_database", + "current_schema", + "current_catalog", + ], + &Inputs::Any, + string, + ); + out.extend(scalar(&["input_file_name"], &Inputs::Any, string)); + out.extend(scalar( + &["input_file_block_start", "input_file_block_length"], + &Inputs::Any, + long, + )); + out.extend(scalar(&["monotonically_increasing_id"], &Inputs::Any, long)); + out.extend(scalar(&["spark_partition_id"], &Inputs::Any, int)); + out.extend(scalar(&["typeof"], &Inputs::Any, typeof_result)); + out.extend(scalar(&["rand", "random"], &Inputs::Any, rand)); + out.extend(scalar(&["randn"], &Inputs::Any, randn)); + out.extend(scalar(&["uuid"], &Inputs::Any, uuid)); + out.extend(scalar(&["uniform"], &Inputs::Any, uniform)); + out.extend(scalar(&["randstr"], &Inputs::Any, randstr)); + out +} diff --git a/crates/krishiv-functions/src/registry/mod.rs b/crates/krishiv-functions/src/registry/mod.rs index c6af46e3e..61f5f73e3 100644 --- a/crates/krishiv-functions/src/registry/mod.rs +++ b/crates/krishiv-functions/src/registry/mod.rs @@ -19,6 +19,7 @@ mod conditional; mod defs; mod hashing; mod math; +mod misc; mod number; mod predicate; diff --git a/crates/krishiv-functions/src/registry/tests.rs b/crates/krishiv-functions/src/registry/tests.rs index 6388993b5..18c766401 100644 --- a/crates/krishiv-functions/src/registry/tests.rs +++ b/crates/krishiv-functions/src/registry/tests.rs @@ -81,15 +81,32 @@ fn every_scalar_definition_binds_and_every_binding_is_defined() { fn lookups_are_case_insensitive_and_versioned() { let r = builtins(); assert!(matches!(r.lookup("UpPeR", SparkProfile::V3_5), Lookup::Found(d) if d.name == "upper")); - // `randstr` arrived in 4.0 and is not implemented. + // `randstr` arrived in 4.0 and is implemented. assert!(matches!( r.lookup("randstr", SparkProfile::V3_5), Lookup::NotFound )); assert!(matches!( r.lookup("RANDSTR", SparkProfile::V4_0), - Lookup::KnownUnsupported + Lookup::Found(_) )); + // So are `uniform` (4.0) and the misc functions. + for name in [ + "uniform", + "uuid", + "typeof", + "version", + "current_database", + "current_catalog", + "monotonically_increasing_id", + "spark_partition_id", + "input_file_name", + ] { + assert!( + matches!(r.lookup(name, SparkProfile::V4_2), Lookup::Found(_)), + "{name}" + ); + } assert!(matches!( r.lookup("no_such_fn", SparkProfile::V4_2), Lookup::NotFound @@ -132,14 +149,22 @@ fn kinds_determinism_and_execution() { assert_eq!(d.execution, Execution::Evaluator, "{name}"); assert_eq!(d.kind, FunctionKind::Scalar, "{name}"); } - for name in ["rand", "randn", "uuid", "shuffle"] { - // Not implemented here, but the optimizer's flag comes from Spark. + for name in ["rand", "randn", "uuid", "uniform", "randstr", "shuffle"] { + // The optimizer's flag comes from Spark; the implemented ones are + // nondeterministic scalar bindings (`shuffle` stays unimplemented: + // it needs an array type the engine has no representation for). assert_eq!( builtins().spark_function(name).unwrap().deterministic, Some(false), "{name}" ); } + for name in ["rand", "randn", "uuid", "uniform", "randstr"] { + let d = def(name); + assert!(!d.deterministic, "{name}"); + assert_eq!(d.execution, Execution::Scalar, "{name}"); + assert_eq!(d.kind, FunctionKind::Scalar, "{name}"); + } for name in [ "round", "pmod", "hash", "try_add", "greatest", "abs", "sha2", ] { @@ -711,6 +736,251 @@ fn session_and_clock_functions() { builtins().lookup("current_time", SparkProfile::V4_2), Lookup::KnownUnsupported )); + // `shuffle` needs an array type the engine has no representation for: + // it stays known-unsupported where Spark has it. + assert!(matches!( + builtins().lookup("shuffle", SparkProfile::V4_2), + Lookup::KnownUnsupported + )); +} + +/// The misc and seeded-random signatures: arities, inputs and result types. +#[test] +fn misc_signatures() { + let (s, i, l) = ( + ty(DataType::String), + ty(DataType::Integer), + ty(DataType::Long), + ); + // The session constants and `version` take no arguments and are never + // NULL. + for name in [ + "version", + "current_database", + "current_schema", + "current_catalog", + "input_file_name", + "monotonically_increasing_id", + "spark_partition_id", + "typeof", + "uuid", + "rand", + "random", + "randn", + ] { + let d = def(name); + assert_eq!(d.kind, FunctionKind::Scalar, "{name}"); + assert_eq!(d.execution, Execution::Scalar, "{name}"); + } + for name in [ + "version", + "current_database", + "current_schema", + "current_catalog", + ] { + assert_eq!(result(name, &[], &[]).unwrap(), s, "{name}"); + assert!(def(name).deterministic, "{name}"); + } + // The file placeholders and partition functions are nondeterministic. + for name in [ + "input_file_name", + "input_file_block_start", + "input_file_block_length", + "monotonically_increasing_id", + "spark_partition_id", + ] { + assert!(!def(name).deterministic, "{name}"); + } + assert_eq!(result("input_file_name", &[], &[]).unwrap(), s); + assert_eq!( + result("input_file_block_start", &[], &[]).unwrap(), + ty(DataType::Long) + ); + assert_eq!( + result("input_file_block_length", &[], &[]).unwrap(), + ty(DataType::Long) + ); + assert_eq!(result("monotonically_increasing_id", &[], &[]).unwrap(), l); + assert_eq!(result("spark_partition_id", &[], &[]).unwrap(), i); + assert_eq!(result("typeof", std::slice::from_ref(&s), &[]).unwrap(), s); + assert_eq!(result("uuid", &[], &[]).unwrap(), s); + assert_eq!(result("rand", &[], &[]).unwrap(), ty(DataType::Double)); + assert_eq!( + result( + "rand", + std::slice::from_ref(&l), + &[Some(Constant::Integral(0))] + ) + .unwrap(), + ty(DataType::Double) + ); + // A NULL seed reads as 0. + assert_eq!( + result("rand", &[null_ty()], &[Some(Constant::Null)]).unwrap(), + ty(DataType::Double) + ); + // A short seed is the wrong type; a non-foldable one is refused. + let err = result( + "rand", + &[ty(DataType::Short)], + &[Some(Constant::Integral(1))], + ) + .unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(err.to_string().contains("(\"INT\" or \"BIGINT\")"), "{err}"); + let err = result("rand", std::slice::from_ref(&i), &[]).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::SeedExpressionIsUnfoldable); + assert!(err.to_string().contains("must be foldable"), "{err}"); + // `uuid` wants a plain literal. + let err = result("uuid", std::slice::from_ref(&i), &[]).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::InvalidParameterValueLong); + assert!(err.to_string().contains("`UUID`"), "{err}"); + // A non-integral literal is the wrong seed too. + let err = result( + "uuid", + &[ty(DataType::Double)], + &[Some(Constant::Fractional(1))], + ) + .unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::InvalidParameterValueLong); +} + +/// `uniform` and `randstr`: the golden types and errors. +#[test] +fn uniform_and_randstr_signatures() { + let (i, l) = (ty(DataType::Integer), ty(DataType::Long)); + let (ten, twenty, zero) = ( + Some(Constant::Integral(10)), + Some(Constant::Integral(20)), + Some(Constant::Integral(0)), + ); + // Integral bounds widen; the seed does not change the type. + assert_eq!( + result( + "uniform", + &[i.clone(), i.clone(), i.clone()], + &[ten.clone(), twenty.clone(), zero.clone()] + ) + .unwrap(), + i + ); + assert_eq!( + result( + "uniform", + &[i.clone(), l.clone(), l.clone()], + &[ten.clone(), twenty.clone(), zero.clone()] + ) + .unwrap(), + l + ); + // A NULL bound is a NULL double; a NULL seed reads as 0. + assert_eq!( + result( + "uniform", + &[null_ty(), i.clone(), i.clone()], + &[Some(Constant::Null), twenty.clone(), zero.clone()] + ) + .unwrap(), + ArgType::new(DataType::Double, true) + ); + // Narrower integrals widen; decimals meet at the wider. + assert_eq!( + result( + "uniform", + &[ty(DataType::Byte), ty(DataType::Short)], + &[Some(Constant::Integral(1)), Some(Constant::Integral(2))] + ) + .unwrap(), + ty(DataType::Short) + ); + assert_eq!( + result( + "uniform", + &[ty(dec(10, 3)), ty(dec(11, 4))], + &[Some(Constant::Fractional(1)), Some(Constant::Fractional(2))] + ) + .unwrap(), + ty(dec(11, 4)) + ); + assert_eq!( + result( + "uniform", + &[ty(DataType::Float), ty(DataType::Double)], + &[Some(Constant::Fractional(1)), Some(Constant::Fractional(2))] + ) + .unwrap(), + ty(DataType::Double) + ); + assert_eq!( + result( + "uniform", + &[i.clone(), i.clone(), null_ty()], + &[ten.clone(), twenty.clone(), Some(Constant::Null)] + ) + .unwrap(), + ArgType::new(DataType::Integer, true) + ); + // A string bound is the wrong type; a column seed is not foldable. + let s = ty(DataType::String); + let err = result( + "uniform", + &[s, i.clone(), i.clone()], + &[None, twenty.clone(), zero.clone()], + ) + .unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(err.to_string().contains("\"NUMERIC\""), "{err}"); + let err = result( + "uniform", + &[i.clone(), i.clone(), i.clone()], + &[ten.clone(), twenty.clone(), None], + ) + .unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::DatatypeMismatchNonFoldableInput + ); + assert!(err.to_string().contains("`seed`"), "{err}"); + // A failing fold propagates instead of typing the call. + let failed = Constant::Error(KrishivError::from(krishiv_common::error::InternalError { + message: "boom".to_owned(), + })); + let err = result( + "uniform", + &[i.clone(), i.clone(), i.clone()], + &[ten.clone(), twenty, Some(failed.clone())], + ) + .unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::InternalError); + let err = result( + "randstr", + &[i.clone(), i.clone()], + &[ten.clone(), Some(failed)], + ) + .unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::InternalError); + // `randstr` is a string; a boolean length is the wrong type. + assert_eq!( + result("randstr", &[i.clone(), i.clone()], &[ten, zero]).unwrap(), + ty(DataType::String) + ); + let err = result( + "randstr", + &[ty(DataType::Boolean), i], + &[Some(Constant::Boolean(true)), None], + ) + .unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(err.to_string().contains("\"INT\""), "{err}"); } #[test] diff --git a/crates/krishiv-functions/src/scalar/misc.rs b/crates/krishiv-functions/src/scalar/misc.rs new file mode 100644 index 000000000..ef7d3ce42 --- /dev/null +++ b/crates/krishiv-functions/src/scalar/misc.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Misc scalar functions: `version`, `current_database` (`current_schema`), +//! `current_catalog`, `input_file_name` (and the block functions), +//! `monotonically_increasing_id` and `spark_partition_id`. +//! +//! Ground truth: Spark 4.2's `misc.scala` (`SparkVersion`, `CurrentDatabase`, +//! `CurrentCatalog`, `InputFileName`, ...), `MonotonicallyIncreasingID.scala` +//! and `SparkPartitionID.scala`, plus the `current_database_catalog.sql` +//! golden. +//! +//! `version` reports the kernels' Spark line (`4.2.0` for the default +//! build). The database and catalog are the engine's single namespace: +//! `default` in `spark_catalog` (see `krishiv-session`'s catalog), so they +//! bind as constants, with no session state. `input_file_name` and the +//! block functions read the file being scanned; the evaluator has no file +//! context yet, so they return what Spark returns with no file being read +//! (the empty string, `-1`). +//! +//! `monotonically_increasing_id` is `(partition << 33) + row_in_partition` +//! and `spark_partition_id` is the partition index, both from the +//! executing partition (see [`super::ScalarFunction::invoke_in`]); each +//! bound function counts its partition's rows in the partitioned body. + +use std::sync::{Arc, Mutex}; + +use arrow_array::{ArrayRef, Int32Array, Int64Array, StringArray}; +use arrow_schema::DataType; +use krishiv_common::error::KrishivError; +use krishiv_kernels::SparkVersion; + +use super::args::{Inputs, internal, unsupported}; +use super::{BindArg, BindOptions, Binding, FamilyResult, bad_types, listed}; + +/// Every name this family binds. `typeof` is resolved by the analyzer +/// (a literal rewrite, like Spark's `RuntimeReplaceable`); its binding is +/// unreachable and reports that. +pub(super) const NAMES: &[&str] = &[ + "version", + "current_database", + "current_schema", + "current_catalog", + "input_file_name", + "input_file_block_start", + "input_file_block_length", + "monotonically_increasing_id", + "spark_partition_id", + "typeof", +]; + +/// The engine's single namespace and catalog (the only values +/// `current_database` / `current_catalog` can report). +const DATABASE: &str = "default"; +const CATALOG: &str = "spark_catalog"; + +/// The release `version` reports for `kernel`. +fn release(kernel: SparkVersion) -> Result<&'static str, KrishivError> { + match kernel { + SparkVersion::V4_0 => Ok("4.0.0"), + SparkVersion::V4_1 => Ok("4.1.0"), + SparkVersion::V4_2 => Ok("4.2.0"), + _ => Err(internal( + "a Spark version with no release string".to_owned(), + )), + } +} + +/// `version()`: the kernels' release line. +fn version(opts: &BindOptions) -> Result { + let version = Arc::::from(release(opts.kernel.spark_version())?); + Ok(Binding::new(DataType::Utf8, move |x: &Inputs<'_>| { + Ok(Arc::new(StringArray::from(vec![&*version; x.rows()])) as ArrayRef) + })) +} + +/// `current_database()` / `current_schema()` / `current_catalog()`. +fn catalog(name: &'static str) -> Result { + let value = Arc::::from(match name { + "current_catalog" => CATALOG, + _ => DATABASE, + }); + Ok(Binding::new(DataType::Utf8, move |x: &Inputs<'_>| { + Ok(Arc::new(StringArray::from(vec![&*value; x.rows()])) as ArrayRef) + })) +} + +/// `input_file_name()`: the empty string (no file is being read). +fn input_file_name() -> Binding { + Binding::new(DataType::Utf8, |x: &Inputs<'_>| { + Ok(Arc::new(StringArray::from(vec![""; x.rows()])) as ArrayRef) + }) +} + +/// `input_file_block_start()` / `input_file_block_length()`: -1 (no file +/// is being read). +fn input_file_block() -> Binding { + Binding::new(DataType::Int64, move |x: &Inputs<'_>| { + Ok(Arc::new(Int64Array::from(vec![-1; x.rows()])) as ArrayRef) + }) +} + +/// `spark_partition_id()`: the executing partition's index. +fn spark_partition_id() -> Binding { + Binding::partitioned(DataType::Int32, |x: &Inputs<'_>, partition| { + Ok(Arc::new(Int32Array::from(vec![partition as i32; x.rows()])) as ArrayRef) + }) +} + +/// `monotonically_increasing_id()`: `(partition << 33) + row_in_partition`, +/// counted per partition in the partitioned body. +fn monotonically_increasing_id() -> Binding { + let counts = Mutex::new(Vec::::new()); + Binding::partitioned(DataType::Int64, move |x: &Inputs<'_>, partition| { + let mut counts = counts + .lock() + .map_err(|_| internal("a row counter is poisoned".to_owned()))?; + let index = partition as usize; + if counts.len() <= index { + counts.resize(index + 1, 0); + } + let first = counts[index]; + // Java's `long` arithmetic wraps; partitions past `2^31` wrap the + // same way instead of panicking. + let base = (u64::from(partition) << 33).wrapping_add(first); + let rows = x.rows() as u64; + counts[index] = first.wrapping_add(rows); + Ok(Arc::new(Int64Array::from( + (0..rows) + .map(|i| base.wrapping_add(i) as i64) + .collect::>(), + )) as ArrayRef) + }) +} + +/// Binds the misc function `name`. +pub(super) fn bind(name: &str, args: &[BindArg], opts: &BindOptions) -> FamilyResult { + let listed = listed(NAMES, name)?; + if !args.is_empty() { + return Some((listed, Err(bad_types(listed, args)))); + } + // `typeof` never reaches execution: the analyzer rewrites it to a + // string literal, so a call that arrives here is an analyzer bug. + if listed == "typeof" { + return Some(( + listed, + Err(unsupported( + "typeof without the analyzer's literal rewrite".to_owned(), + )), + )); + } + let binding = match listed { + "version" => version(opts), + "current_database" | "current_schema" | "current_catalog" => catalog(listed), + "input_file_name" => Ok(input_file_name()), + "input_file_block_start" | "input_file_block_length" => Ok(input_file_block()), + "spark_partition_id" => Ok(spark_partition_id()), + "monotonically_increasing_id" => Ok(monotonically_increasing_id()), + _ => return Some((listed, Err(bad_types(listed, args)))), + }; + Some((listed, binding)) +} diff --git a/crates/krishiv-functions/src/scalar/mod.rs b/crates/krishiv-functions/src/scalar/mod.rs index def9ab20c..84939f528 100644 --- a/crates/krishiv-functions/src/scalar/mod.rs +++ b/crates/krishiv-functions/src/scalar/mod.rs @@ -29,8 +29,10 @@ mod bitwise; mod datetime; mod hash; mod math; +mod misc; mod number; mod predicate; +mod random; mod regex; mod session; mod string; @@ -39,6 +41,7 @@ pub use regex::{extract_all_pieces, split_regex}; use std::fmt; use std::sync::Arc; +use std::sync::atomic::AtomicU64; use arrow_array::ArrayRef; use arrow_schema::DataType; @@ -99,6 +102,11 @@ pub struct BindOptions { pub decimal_precision_loss: bool, /// The session's user (`current_user()`). pub user: Arc, + /// Counts the unseeded random binds of one plan instantiation, in + /// compilation order: each draws its own query-seeded stream (Spark's + /// `ResolveRandomSeed` assigns each its own seed), while one plan + /// always draws the same streams. + pub seed_salt: Arc, } impl BindOptions { @@ -113,6 +121,7 @@ impl BindOptions { query_start_micros: 0, decimal_precision_loss: true, user: Arc::from("unknown"), + seed_salt: Arc::new(AtomicU64::new(0)), } } @@ -151,6 +160,14 @@ impl BindOptions { self } + /// With this counter of unseeded random binds (see + /// [`Self::seed_salt`]). + #[must_use] + pub fn with_seed_salt(mut self, seed_salt: Arc) -> Self { + self.seed_salt = seed_salt; + self + } + /// With this query start instant (microseconds since the epoch, UTC). #[must_use] pub fn with_query_start(mut self, micros: i64) -> Self { @@ -180,18 +197,45 @@ pub trait ScalarFunction: Send + Sync + fmt::Debug { /// The function's Spark errors (`ARITHMETIC_OVERFLOW` under ANSI, ...), /// and `INTERNAL_ERROR` for arguments that do not match the binding. fn invoke(&self, args: &[Arg], rows: usize) -> Result; + /// As [`Self::invoke`], in the executing partition `partition` (Spark's + /// partition index): a stateful function (`rand`, `uuid`, + /// `monotonically_increasing_id`, ...) owns one value stream per + /// partition, and this evaluates the next `rows` values of that + /// stream. The default runs [`Self::invoke`], for the stateless + /// functions. + /// + /// # Errors + /// As [`Self::invoke`]. + fn invoke_in( + &self, + args: &[Arg], + rows: usize, + partition: u32, + ) -> Result { + let _ = partition; + self.invoke(args, rows) + } } /// The body of a bound function: computes over [`Inputs`] that hold at /// least one column when `rows > 1`. type Body = Box) -> Result + Send + Sync>; +/// The body of a stateful bound function: computes over [`Inputs`] in the +/// executing partition `partition`, advancing that partition's stream once +/// per row it produces. A stateful body always computes every row (never +/// the one-row broadcast of [`Body`]), so the stream advances with the +/// rows even when every argument is a literal. +type PartitionedBody = + Box, u32) -> Result + Send + Sync>; + /// The one [`ScalarFunction`] implementation: a name, a type and a body. struct Bound { name: &'static str, types: Vec, return_type: DataType, body: Body, + partitioned: Option, } impl fmt::Debug for Bound { @@ -214,6 +258,23 @@ impl ScalarFunction for Bound { } fn invoke(&self, args: &[Arg], rows: usize) -> Result { + self.run(args, rows, 0) + } + + fn invoke_in( + &self, + args: &[Arg], + rows: usize, + partition: u32, + ) -> Result { + self.run(args, rows, partition) + } +} + +impl Bound { + /// Checks `args` against the binding, then computes `rows` values in + /// `partition`. + fn run(&self, args: &[Arg], rows: usize, partition: u32) -> Result { if args.len() != self.types.len() { return Err(internal(format!( "{} bound to {} arguments called with {}", @@ -229,9 +290,13 @@ impl ScalarFunction for Bound { if rows == 0 { return Ok(arrow_array::new_empty_array(&self.return_type)); } - // Every argument a literal: compute one row, then repeat it. - let all_literal = args.iter().all(|a| matches!(a, Arg::Scalar(_))); - let out = if all_literal { + let out = if let Some(run) = &self.partitioned { + // A stateful body always computes every row, in partition + // `partition`: the one-row broadcast below would repeat one + // value instead of advancing the partition's stream. + run(&Inputs::new(args, rows), partition)? + } else if args.iter().all(|a| matches!(a, Arg::Scalar(_))) { + // Every argument a literal: compute one row, then repeat it. let one: Vec = args.iter().map(Arg::as_column).collect(); let row = (self.body)(&Inputs::new(&one, 1))?; if rows == 1 { @@ -260,6 +325,7 @@ impl ScalarFunction for Bound { struct Binding { return_type: DataType, body: Body, + partitioned: Option, } impl Binding { @@ -270,6 +336,24 @@ impl Binding { Self { return_type, body: Box::new(body), + partitioned: None, + } + } + + /// A stateful binding: `body` computes every row in the executing + /// partition (see [`PartitionedBody`]). + fn partitioned( + return_type: DataType, + body: impl Fn(&Inputs<'_>, u32) -> Result + Send + Sync + 'static, + ) -> Self { + Self { + return_type, + body: Box::new(|_: &Inputs<'_>| { + Err(internal( + "a stateful function evaluated without a partition".to_owned(), + )) + }), + partitioned: Some(Box::new(body)), } } } @@ -283,7 +367,7 @@ type FamilyResult = Option<(&'static str, Result)>; type Binder = fn(&str, &[BindArg], &BindOptions) -> FamilyResult; /// The families (their names and binder), in lookup order. -const FAMILIES: [(&[&str], Binder); 10] = [ +const FAMILIES: [(&[&str], Binder); 12] = [ (string::NAMES, string::bind), (datetime::NAMES, datetime::bind), (math::NAMES, math::bind), @@ -294,6 +378,8 @@ const FAMILIES: [(&[&str], Binder); 10] = [ (bitwise::NAMES, bitwise::bind), (predicate::NAMES, predicate::bind), (number::NAMES, number::bind), + (misc::NAMES, misc::bind), + (random::NAMES, random::bind), ]; /// The name of every scalar function [`bind_scalar`] binds, aliases @@ -355,6 +441,7 @@ pub fn bind_scalar( types: args.iter().map(BindArg::data_type).collect(), return_type: binding.return_type, body: binding.body, + partitioned: binding.partitioned, })) } diff --git a/crates/krishiv-functions/src/scalar/random.rs b/crates/krishiv-functions/src/scalar/random.rs new file mode 100644 index 000000000..0495120ae --- /dev/null +++ b/crates/krishiv-functions/src/scalar/random.rs @@ -0,0 +1,533 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Seeded random functions: `rand` (`random`), `randn`, `uniform`, +//! `randstr` and `uuid`. +//! +//! Ground truth: Spark 4.2's `randomExpressions.scala` (`Rand`, `Randn`, +//! `Uniform`, `RandStr`), `misc.scala` (`Uuid`) and +//! `XORShiftRandom` (`core/.../util/random/XORShiftRandom.scala`), plus the +//! `random.sql` golden (`tests/spark-golden/v4.2.0/results/random.sql.out`). +//! +//! Every function owns one value stream per partition: the stream is a pure +//! function of the seed (the folded seed argument, or the query seed for an +//! absent or NULL seed) and the partition index, advanced once per row the +//! partition produces. The analyzer only lets foldable seeds through (the +//! registry rules), so a bind reads the seed from the first row of every +//! invocation; the per-partition streams live in the partitioned body. +//! +//! `rand`/`randn` draw from Spark's `XORShiftRandom` (the MurmurHash3 seed +//! hash and the shift register below), so seeded results match Spark +//! exactly (`rand(0)` is `0.7604953758285915`). `uuid` draws its 128 bits +//! from the same generator and sets the RFC 4122 version and variant bits; +//! Spark fills those bits from a Mersenne Twister instead, so only the +//! shape (and the per-seed determinism) matches, not Spark's values. + +use std::sync::{Arc, Mutex}; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float64Type, Int32Type, Int64Type}; +use arrow_array::{ArrayRef, Float64Array, StringArray}; +use arrow_schema::DataType; +use krishiv_common::error::KrishivError; +use krishiv_kernels::cast::{CastMode, CastOptions, TimeZoneRules, cast}; + +use super::args::{Inputs, internal}; +use super::{BindArg, BindOptions, Binding, FamilyResult, bad_types, listed}; + +/// Every name this family binds (`random` is Spark's alias of `rand`). +pub(super) const NAMES: &[&str] = &["rand", "random", "randn", "uniform", "randstr", "uuid"]; + +/// Spark's `XORShiftRandom`: the seed hash (Scala's `MurmurHash3`, the only +/// hashing in the stream) and the shift register. +#[derive(Debug, Clone)] +struct XorShift { + seed: u64, +} + +impl XorShift { + /// The stream for `init` (Spark's `new XORShiftRandom(init)`). + fn new(init: i64) -> Self { + Self { + seed: hash_seed(init), + } + } + + /// Spark's `next(bits)`. + #[expect( + clippy::cast_possible_truncation, + reason = "Java's `(int)` cast of the masked bits" + )] + fn next(&mut self, bits: u32) -> i32 { + let mut next = self.seed ^ self.seed.wrapping_shl(21); + next ^= next >> 35; + next ^= next.wrapping_shl(4); + self.seed = next; + (next & ((1_u64 << bits) - 1)) as i32 + } + + /// `java.util.Random.nextDouble`: 53 bits over `2^53`. + fn next_double(&mut self) -> f64 { + let high = i64::from(self.next(26)); + let low = i64::from(self.next(27)); + ((high << 27) + low) as f64 / 9_007_199_254_740_992.0 + } + + /// `java.util.Random.nextGaussian` (the polar method): `cached` is the + /// spare value (`haveNextNextGaussian` / `nextNextGaussian`). + fn next_gaussian(&mut self, cached: &mut Option) -> f64 { + if let Some(v) = cached.take() { + return v; + } + loop { + let v1 = 2.0 * self.next_double() - 1.0; + let v2 = 2.0 * self.next_double() - 1.0; + let s = v1 * v1 + v2 * v2; + if s < 1.0 && s != 0.0 { + let multiplier = (-2.0 * s.ln() / s).sqrt(); + *cached = Some(v2 * multiplier); + return v1 * multiplier; + } + } + } + + /// `java.util.Random.nextLong`. + fn next_long(&mut self) -> i64 { + (i64::from(self.next(32)) << 32) + i64::from(self.next(32)) + } +} + +/// Spark's `XORShiftRandom.hashSeed`: the seed's big-endian bytes hashed +/// with Scala's `MurmurHash3.bytesHash`. +fn hash_seed(seed: i64) -> u64 { + let bytes = seed.to_be_bytes(); + let low = bytes_hash(&bytes, 0x3c07_4a61); + let high = bytes_hash(&bytes, low); + ((i64::from(high) << 32) | (i64::from(low) & 0xffff_ffff)).cast_unsigned() +} + +/// Scala 2.13's `MurmurHash3.bytesHash` (an 8-byte input has no tail). +fn bytes_hash(data: &[u8], seed: i32) -> i32 { + let mut h = seed; + let (chunks, _) = data.as_chunks::<4>(); + for chunk in chunks { + let k = i32::from_le_bytes(*chunk); + h = murmur_mix(h, k); + } + murmur_finalize(h, data.len()) +} + +/// Scala 2.13's `MurmurHash3.mix`. +fn murmur_mix(hash: i32, data: i32) -> i32 { + murmur_mix_last(hash, data) + .rotate_left(13) + .wrapping_mul(5) + .wrapping_add(0xe654_6b64u32.cast_signed()) +} + +/// Scala 2.13's `MurmurHash3.mixLast`. +fn murmur_mix_last(hash: i32, mut data: i32) -> i32 { + data = data.wrapping_mul(0xcc9e_2d51u32.cast_signed()); + data = data.rotate_left(15); + data = data.wrapping_mul(0x1b87_3593u32.cast_signed()); + hash ^ data +} + +/// Scala 2.13's `MurmurHash3.finalizeHash`. +#[expect( + clippy::cast_possible_truncation, + reason = "the hashed seed bytes always fit" +)] +fn murmur_finalize(hash: i32, length: usize) -> i32 { + murmur_avalanche(hash ^ length as i32) +} + +/// Scala 2.13's `MurmurHash3.avalanche`. +fn murmur_avalanche(mut hash: i32) -> i32 { + hash ^= (hash.cast_unsigned() >> 16).cast_signed(); + hash = hash.wrapping_mul(0x85eb_ca6bu32.cast_signed()); + hash ^= (hash.cast_unsigned() >> 13).cast_signed(); + hash = hash.wrapping_mul(0xc2b2_ae35u32.cast_signed()); + hash ^= (hash.cast_unsigned() >> 16).cast_signed(); + hash +} + +/// SplitMix64: the per-query seed from the query start instant (fixed per +/// plan, never sampled at execution). +fn splitmix64(mut state: u64) -> u64 { + state = state.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut z = state; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ (z >> 31) +} + +/// The query seed: one value per query, from the plan-fixed query start. +fn query_seed(opts: &BindOptions) -> u64 { + opts.query_start_micros.cast_unsigned() +} + +/// An unseeded (or NULL-seeded, where Spark draws randomly) call's seed: +/// the query seed mixed with a per-occurrence salt, so two unseeded calls +/// in one query draw different streams (Spark's `ResolveRandomSeed` +/// assigns each its own random seed) while one query always draws the same +/// streams. The salt counts the unseeded binds of one plan instantiation, +/// in compilation order. +fn unseeded_seed(opts: &BindOptions) -> i64 { + use std::sync::atomic::Ordering; + let salt = opts.seed_salt.fetch_add(1, Ordering::Relaxed); + splitmix64(query_seed(opts).wrapping_add(salt)) as i64 +} + +/// One value stream: the generator and the spare Gaussian. +#[derive(Debug)] +struct Stream { + rng: XorShift, + gaussian: Option, +} + +impl Stream { + fn new(seed: i64) -> Self { + Self { + rng: XorShift::new(seed), + gaussian: None, + } + } +} + +/// The per-partition streams of one bound function, by partition index. +#[derive(Debug, Default)] +struct Streams { + streams: Vec, +} + +impl Streams { + /// The stream of `partition`, seeded once. + fn stream(&mut self, partition: u32, seed: i64) -> &mut Stream { + let index = partition as usize; + if self.streams.len() <= index { + self.streams.resize_with(index + 1, || Stream::new(seed)); + } + &mut self.streams[index] + } +} + +/// Locks `mutex` (poison is impossible: the bodies never panic). +fn lock(mutex: &Mutex) -> Result, KrishivError> { + mutex + .lock() + .map_err(|_| internal("a random function's streams are poisoned".to_owned())) +} + +/// The folded seed argument `i` (`None` for a NULL seed). +fn seed(x: &Inputs<'_>, i: usize) -> Result, KrishivError> { + let c = x.column(i)?; + if c.data_type() == &DataType::Null || c.is_null(0) { + return Ok(None); + } + match c.data_type() { + DataType::Int32 => Ok(Some(i64::from(c.as_primitive::().value(0)))), + DataType::Int64 => Ok(Some(c.as_primitive::().value(0))), + t => Err(internal(format!("a random seed of type {t}"))), + } +} + +/// The cast to a bound of `uniform` (Spark's plain `Cast`, never `TRY_CAST`). +fn bound_cast( + c: &ArrayRef, + to: &DataType, + ansi: bool, + zone: &Arc, +) -> Result { + cast( + c.as_ref(), + to, + &CastOptions { + mode: if ansi { + CastMode::Ansi + } else { + CastMode::Legacy + }, + time_zone: zone.as_ref(), + }, + ) + .map_err(KrishivError::from) +} + +/// The folded numeric argument `i` as an `f64` (`None` for NULL): Spark's +/// `Cast` to `DOUBLE`. +fn double( + x: &Inputs<'_>, + i: usize, + ansi: bool, + zone: &Arc, +) -> Result, KrishivError> { + let c = x.column(i)?; + if c.data_type() == &DataType::Null || c.is_null(0) { + return Ok(None); + } + let d = bound_cast(&c, &DataType::Float64, ansi, zone)?; + if d.is_null(0) { + return Ok(None); + } + Ok(Some(d.as_primitive::().value(0))) +} + +/// The folded length argument `i` as an `i32` (`None` for NULL): Spark's +/// `Cast` to `INT`. +fn length( + x: &Inputs<'_>, + i: usize, + ansi: bool, + zone: &Arc, +) -> Result, KrishivError> { + let c = x.column(i)?; + if c.data_type() == &DataType::Null || c.is_null(0) { + return Ok(None); + } + let d = bound_cast(&c, &DataType::Int32, ansi, zone)?; + if d.is_null(0) { + return Ok(None); + } + Ok(Some(d.as_primitive::().value(0))) +} + +/// `rand(seed?)` / `random(seed?)`: `rng.nextDouble()` per row. +fn rand(opts: &BindOptions, args: &[BindArg]) -> Binding { + let queries = unseeded_seed(opts); + let seeded = !args.is_empty(); + let state = Mutex::new(Streams::default()); + Binding::partitioned(DataType::Float64, move |x, partition| { + let seed = if seeded { + seed(x, 0)?.unwrap_or(0) + } else { + queries + }; + let mut streams = lock(&state)?; + let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let out: Vec = (0..x.rows()).map(|_| s.rng.next_double()).collect(); + Ok(Arc::new(Float64Array::from(out)) as ArrayRef) + }) +} + +/// `randn(seed?)`: `rng.nextGaussian()` per row. +fn randn(opts: &BindOptions, args: &[BindArg]) -> Binding { + let queries = unseeded_seed(opts); + let seeded = !args.is_empty(); + let state = Mutex::new(Streams::default()); + Binding::partitioned(DataType::Float64, move |x, partition| { + let seed = if seeded { + seed(x, 0)?.unwrap_or(0) + } else { + queries + }; + let mut streams = lock(&state)?; + let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let out: Vec = (0..x.rows()) + .map(|_| s.rng.next_gaussian(&mut s.gaussian)) + .collect(); + Ok(Arc::new(Float64Array::from(out)) as ArrayRef) + }) +} + +/// `uuid(seed?)`: 128 bits per row, with the RFC 4122 version and variant +/// bits. A NULL or absent seed draws from the query seed. +fn uuid(opts: &BindOptions, args: &[BindArg]) -> Binding { + let queries = unseeded_seed(opts); + let seeded = !args.is_empty(); + let state = Mutex::new(Streams::default()); + Binding::partitioned(DataType::Utf8, move |x, partition| { + let seed = if seeded { + seed(x, 0)?.unwrap_or(queries) + } else { + queries + }; + let mut streams = lock(&state)?; + let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let out: Vec = (0..x.rows()) + .map(|_| { + let most = (s.rng.next_long().cast_unsigned() & 0xffff_ffff_ffff_0fff) + | 0x0000_0000_0000_4000; + let least = (s.rng.next_long().cast_unsigned() | 0x8000_0000_0000_0000) + & 0xbfff_ffff_ffff_ffff; + format!( + "{:08x}-{:04x}-{:04x}-{:04x}-{:012x}", + (most >> 32) as u32, + ((most >> 16) & 0xffff) as u16, + (most & 0xffff) as u16, + ((least >> 48) & 0xffff) as u16, + least & 0xffff_ffff_ffff, + ) + }) + .collect(); + Ok(Arc::new(StringArray::from(out)) as ArrayRef) + }) +} + +/// Whether `t` is an integral Arrow type. +fn is_integral(t: &DataType) -> bool { + matches!( + t, + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 + ) +} + +/// The wider of two integral Arrow types (Spark's `UpCastRule` order). +fn wider_integral(a: &DataType, b: &DataType) -> DataType { + let rank = |t: &DataType| match t { + DataType::Int8 => 0, + DataType::Int16 => 1, + DataType::Int32 => 2, + _ => 3, + }; + if rank(a) >= rank(b) { + a.clone() + } else { + b.clone() + } +} + +/// The Arrow result type of `uniform(min, max)` for argument types +/// `min`/`max` (an untyped `NULL` reads as `DOUBLE`, as Spark casts it): +/// Spark's `Uniform.dataType`. +fn uniform_result(min: &DataType, max: &DataType) -> DataType { + let num = |t: &DataType| { + if *t == DataType::Null { + DataType::Float64 + } else { + t.clone() + } + }; + let (min, max) = (num(min), num(max)); + match (&min, &max) { + (a, b) if is_integral(a) && is_integral(b) => wider_integral(a, b), + (DataType::Float64, _) | (_, DataType::Float64) => DataType::Float64, + (DataType::Float32, _) | (_, DataType::Float32) => DataType::Float32, + (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => { + if is_wider_decimal((*p1, *s1), (*p2, *s2)) { + min + } else { + max + } + } + (_, d @ DataType::Decimal128(..)) | (d @ DataType::Decimal128(..), _) => d.clone(), + _ => DataType::Float64, + } +} + +/// Spark's `DecimalType.isWiderThan`: every part at least as wide. +fn is_wider_decimal(left: (u8, i8), right: (u8, i8)) -> bool { + let (p1, s1) = (i32::from(left.0), i32::from(left.1)); + let (p2, s2) = (i32::from(right.0), i32::from(right.1)); + p1 - s1 >= p2 - s2 && s1 >= s2 && p1 >= p2 +} + +/// `uniform(min, max[, seed])`: Spark's `Uniform` is a `RuntimeReplaceable` +/// over `cast(min + (max - min) * Rand(seed))`, computed here in one pass: +/// the bounds as doubles, one draw per row, then Spark's `Cast` to the +/// result type. A NULL bound is a NULL result. +fn uniform(opts: &BindOptions, args: &[BindArg]) -> Binding { + let types: Vec = args.iter().map(BindArg::data_type).collect(); + let result = uniform_result(&types[0], &types[1]); + let out_type = result.clone(); + let queries = unseeded_seed(opts); + let seeded = args.len() > 2; + let ansi = opts.ansi; + let zone = Arc::clone(&opts.time_zone); + let state = Mutex::new(Streams::default()); + Binding::partitioned(result, move |x, partition| { + let seed = if seeded { + seed(x, 2)?.unwrap_or(0) + } else { + queries + }; + let (Some(lo), Some(hi)) = (double(x, 0, ansi, &zone)?, double(x, 1, ansi, &zone)?) else { + return Ok(arrow_array::new_null_array(&out_type, x.rows())); + }; + let mut streams = lock(&state)?; + let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let draws: Vec = (0..x.rows()) + .map(|_| lo + (hi - lo) * s.rng.next_double()) + .collect(); + let draws = Float64Array::from(draws); + bound_cast(&(Arc::new(draws) as ArrayRef), &out_type, ansi, &zone) + }) +} + +/// `randstr(length[, seed])`: `length` characters per row, each drawn with +/// `Math.abs(rng.nextInt() % 62)` and read as a digit, a lower-case or an +/// upper-case letter (Spark's `ExpressionImplUtils.randStr`). A NULL length +/// is the empty string; a negative length is +/// `INVALID_PARAMETER_VALUE.LENGTH`. +fn randstr(opts: &BindOptions, args: &[BindArg]) -> Binding { + let queries = unseeded_seed(opts); + let seeded = args.len() > 1; + let ansi = opts.ansi; + let zone = Arc::clone(&opts.time_zone); + let state = Mutex::new(Streams::default()); + Binding::partitioned(DataType::Utf8, move |x, partition| { + let seed = if seeded { + seed(x, 1)?.unwrap_or(queries) + } else { + queries + }; + let Some(n) = length(x, 0, ansi, &zone)? else { + return Ok(Arc::new(StringArray::from(vec![""; x.rows()])) as ArrayRef); + }; + if n < 0 { + return Err(krishiv_common::error::InvalidParameterValueLength { + function_name: "`randstr`".to_owned(), + length: n.to_string(), + parameter: "`length`".to_owned(), + } + .into()); + } + let mut streams = lock(&state)?; + let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let out: Vec = (0..x.rows()) + .map(|_| { + (0..n) + .map(|_| { + // Java's `Math.abs(rng.nextInt() % 62)`: `%` keeps + // the dividend's sign, then `abs`. + let v = (s.rng.next(32) % 62).abs(); + // `v` is `abs(nextInt() % 62)`, so in `0..62`. + let v = u8::try_from(v).unwrap_or(0); + if v < 10 { + (b'0' + v) as char + } else if v < 36 { + (b'a' + v - 10) as char + } else { + (b'A' + v - 36) as char + } + }) + .collect() + }) + .collect(); + Ok(Arc::new(StringArray::from(out)) as ArrayRef) + }) +} + +/// Whether `t` is an `int` or `bigint` seed (an untyped `NULL` reads as +/// one): what the analyzer lets through for every seed but `uuid`'s. +fn is_seed(t: &DataType) -> bool { + matches!(t, DataType::Int32 | DataType::Int64 | DataType::Null) +} + +/// Binds the random function `name`. +pub(super) fn bind(name: &str, args: &[BindArg], opts: &BindOptions) -> FamilyResult { + let listed = listed(NAMES, name)?; + let types: Vec = args.iter().map(BindArg::data_type).collect(); + let binding = match (listed, types.as_slice()) { + ("rand" | "random", []) => rand(opts, args), + ("rand" | "random" | "randn" | "uuid", [t]) if is_seed(t) => match listed { + "randn" => randn(opts, args), + "uuid" => uuid(opts, args), + _ => rand(opts, args), + }, + ("randn", []) => randn(opts, args), + ("uuid", []) => uuid(opts, args), + ("uniform", [_, _] | [_, _, _]) => uniform(opts, args), + ("randstr", [_] | [_, _]) => randstr(opts, args), + _ => return Some((listed, Err(bad_types(listed, args)))), + }; + Some((listed, Ok(binding))) +} diff --git a/crates/krishiv-functions/src/scalar/tests/misc.rs b/crates/krishiv-functions/src/scalar/tests/misc.rs new file mode 100644 index 000000000..f57c050bf --- /dev/null +++ b/crates/krishiv-functions/src/scalar/tests/misc.rs @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The misc bindings: the session constants, the file placeholders and the +//! partition functions. + +use std::sync::Arc; + +use arrow_array::Array; +use arrow_array::cast::AsArray; +use arrow_array::types::{Int32Type, Int64Type}; +use arrow_schema::DataType; +use krishiv_common::error::ErrorClass; + +use super::{BindArg, BindOptions, bind, bind_scalar, class, strs}; + +fn strings(f: &Arc, rows: usize, partition: u32) -> Vec { + let out = f.invoke_in(&[], rows, partition).unwrap(); + out.as_string::() + .iter() + .map(|s| s.unwrap().to_owned()) + .collect() +} + +#[test] +fn version_reports_the_kernels_release_line() { + let f = bind("version", &[]); + assert_eq!(f.return_type(), &DataType::Utf8); + assert_eq!(strings(&f, 3, 0), vec!["4.2.0"; 3]); + // The line follows the kernels' Spark version. + let opts = BindOptions::new().with_kernel( + krishiv_kernels::KernelCtx::new().with_spark_version(krishiv_kernels::SparkVersion::V4_0), + ); + let f = super::bind_scalar("version", &[], &opts).unwrap(); + assert_eq!(strings(&f, 1, 0), vec!["4.0.0"]); + // `typeof` never reaches execution: the analyzer rewrites it first. + assert_eq!( + class( + &super::bind_scalar( + "typeof", + &[BindArg::Column(DataType::Int32)], + &BindOptions::new() + ) + .unwrap_err() + ), + ErrorClass::KrishivUnsupportedFeature + ); +} + +#[test] +fn current_database_and_catalog() { + assert_eq!( + strings(&bind("current_database", &[]), 2, 0), + vec!["default"; 2] + ); + assert_eq!(strings(&bind("current_schema", &[]), 1, 0), vec!["default"]); + assert_eq!( + strings(&bind("current_catalog", &[]), 1, 0), + vec!["spark_catalog"] + ); +} + +#[test] +fn input_file_functions_read_no_file() { + assert_eq!(strings(&bind("input_file_name", &[]), 2, 0), vec![""; 2]); + for name in ["input_file_block_start", "input_file_block_length"] { + let f = bind(name, &[]); + assert_eq!(f.return_type(), &DataType::Int64); + let out = f.invoke_in(&[], 3, 0).unwrap(); + assert_eq!( + out.as_primitive::().values(), + &[-1, -1, -1], + "{name}" + ); + } +} + +#[test] +fn spark_partition_id_is_the_partition() { + let f = bind("spark_partition_id", &[]); + assert_eq!(f.return_type(), &DataType::Int32); + for partition in [0, 1, 7] { + let out = f.invoke_in(&[], 3, partition).unwrap(); + assert_eq!( + out.as_primitive::().values(), + &[partition as i32; 3], + "partition {partition}" + ); + } + // Without a partition the single partition reads 0. + let out = f.invoke(&[], 2).unwrap(); + assert_eq!(out.as_primitive::().values(), &[0, 0]); +} + +#[test] +fn monotonically_increasing_id_counts_per_partition() { + let f = bind("monotonically_increasing_id", &[]); + assert_eq!(f.return_type(), &DataType::Int64); + let ids = |rows: usize, partition: u32| { + f.invoke_in(&[], rows, partition) + .unwrap() + .as_primitive::() + .values() + .to_vec() + }; + assert_eq!(ids(3, 0), vec![0, 1, 2]); + // The next batch continues the partition's count. + assert_eq!(ids(2, 0), vec![3, 4]); + // Another partition counts from its own base. + let base = 2_i64 << 33; + assert_eq!(ids(2, 2), vec![base, base + 1]); + assert_eq!(ids(1, 0), vec![5]); +} + +#[test] +fn misc_functions_take_no_arguments() { + for name in [ + "version", + "current_database", + "current_schema", + "current_catalog", + "input_file_name", + "input_file_block_start", + "input_file_block_length", + "spark_partition_id", + "monotonically_increasing_id", + ] { + assert_eq!(bind(name, &[]).name(), name); + let arg = strs(&[Some("x")]); + let err = bind_scalar( + name, + &[BindArg::Column(arg.data_type().clone())], + &BindOptions::new(), + ) + .unwrap_err(); + assert_eq!(class(&err), ErrorClass::KrishivUnsupportedFeature, "{name}"); + // The error names the bad signature, not a missing function. + assert!(err.to_string().contains(name), "{name}: {err}"); + } +} diff --git a/crates/krishiv-functions/src/scalar/tests/mod.rs b/crates/krishiv-functions/src/scalar/tests/mod.rs index 226700c4a..8579b6489 100644 --- a/crates/krishiv-functions/src/scalar/tests/mod.rs +++ b/crates/krishiv-functions/src/scalar/tests/mod.rs @@ -555,6 +555,7 @@ fn results_are_checked_against_the_binding() { types: vec![], return_type: DataType::Int32, body: Box::new(|_| Ok(Arc::new(StringArray::from(vec!["x"])) as ArrayRef)), + partitioned: None, }; let err = wrong.invoke(&[], 1).unwrap_err(); assert_eq!(class(&err), ErrorClass::InternalError); @@ -563,8 +564,10 @@ fn results_are_checked_against_the_binding() { mod bitwise; mod families; +mod misc; mod number; mod predicate; +mod random; #[test] fn session_user_and_local_clock() { diff --git a/crates/krishiv-functions/src/scalar/tests/random.rs b/crates/krishiv-functions/src/scalar/tests/random.rs new file mode 100644 index 000000000..60b3efeb7 --- /dev/null +++ b/crates/krishiv-functions/src/scalar/tests/random.rs @@ -0,0 +1,375 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The seeded random bindings: Spark's streams, golden values and the +//! per-partition contract. + +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float64Type, Int32Type}; +use arrow_array::{Array, ArrayRef, Int32Array, Int64Array}; +use arrow_schema::DataType; + +use super::{Arg, BindArg, BindOptions, ints, lit, strs}; + +fn seed_i32(v: i32) -> ArrayRef { + Arc::new(Int32Array::from(vec![v])) +} + +fn seed_i64(v: i64) -> ArrayRef { + Arc::new(Int64Array::from(vec![v])) +} + +fn null_i32() -> ArrayRef { + Arc::new(Int32Array::from(vec![None as Option])) +} + +/// The `f64` values of `a`. +fn f64s(a: &ArrayRef) -> Vec> { + a.as_primitive::().iter().collect() +} + +/// `f` over `rows` rows in `partition`, with a literal seed `s`. +fn seeded(name: &str, s: &ArrayRef, rows: usize, partition: u32) -> ArrayRef { + let bound = vec![BindArg::Literal(Arc::clone(s))]; + let f = super::bind_scalar(name, &bound, &BindOptions::new()).unwrap(); + assert_eq!(f.return_type(), &DataType::Float64); + f.invoke_in(&[Arg::Scalar(Arc::clone(s))], rows, partition) + .unwrap() +} + +#[test] +fn rand_zero_matches_spark() { + // `SELECT rand(0)` in `random.sql`. + let s = seed_i32(0); + let got = seeded("rand", &s, 1, 0); + assert_eq!(f64s(&got), vec![Some(0.7604953758285915)]); + // The alias binds the same stream. + let got = seeded("random", &s, 1, 0); + assert_eq!(f64s(&got), vec![Some(0.7604953758285915)]); +} + +#[test] +fn rand_advances_one_draw_per_row() { + // `SELECT uniform(10, 20, 0) FROM VALUES (0), (1), (2)` is 15, 16, 17: + // the rows draw consecutively from one stream. + let s = seed_i32(0); + let got = seeded("rand", &s, 3, 0); + let rows = f64s(&got); + assert_eq!(rows.len(), 3); + assert!( + (rows[0].unwrap() - 0.7604953758285915).abs() < 1e-15, + "{rows:?}" + ); + assert!(rows.windows(2).all(|w| w[0] != w[1]), "{rows:?}"); +} + +#[test] +fn randn_zero_matches_spark() { + // `SELECT randn(0)` in `random.sql`. + let s = seed_i64(0); + let got = seeded("randn", &s, 1, 0); + assert_eq!(f64s(&got), vec![Some(1.6034991609278433)]); +} + +#[test] +fn null_seeds_read_as_zero() { + // `rand(NULL)` and `rand(cast(NULL AS INT))` are `rand(0)`. + let s = null_i32(); + let got = seeded("rand", &s, 1, 0); + assert_eq!(f64s(&got), vec![Some(0.7604953758285915)]); + let got = seeded("randn", &s, 1, 0); + assert_eq!(f64s(&got), vec![Some(1.6034991609278433)]); +} + +#[test] +fn long_seeds_match_int_seeds() { + let i = seed_i32(0); + let l = seed_i64(0); + assert_eq!( + f64s(&seeded("rand", &i, 2, 0)), + f64s(&seeded("rand", &l, 2, 0)) + ); + assert_eq!( + f64s(&seeded("randn", &i, 2, 0)), + f64s(&seeded("randn", &l, 2, 0)) + ); +} + +#[test] +fn partitions_own_their_streams() { + // The same query in another partition draws another stream; the + // partition seeds the generator (`seed + partitionIndex`, Spark's + // `initializeInternal`). + let s = seed_i32(0); + let p0 = f64s(&seeded("rand", &s, 3, 0)); + let p1 = f64s(&seeded("rand", &s, 3, 1)); + assert_ne!(p0, p1); + // `monotonically_increasing_id` and `rand` of two partitions differ, + // and one partition is stable across batches. + let f = super::bind_scalar( + "rand", + &[BindArg::Literal(Arc::clone(&s))], + &BindOptions::new(), + ) + .unwrap(); + let lit = Arg::Scalar(Arc::clone(&s)); + let a = f.invoke_in(std::slice::from_ref(&lit), 2, 0).unwrap(); + let b = f.invoke_in(std::slice::from_ref(&lit), 2, 0).unwrap(); + assert_ne!(f64s(&a), f64s(&b)); +} + +#[test] +fn uniform_matches_the_golden_values() { + // `SELECT uniform(0, 10, 0)` is 7, `(10, 20, 0)` is 17. + let bound = |lo: &ArrayRef, hi: &ArrayRef, sd: &ArrayRef| { + let bound = vec![ + BindArg::Literal(Arc::clone(lo)), + BindArg::Literal(Arc::clone(hi)), + BindArg::Literal(Arc::clone(sd)), + ]; + super::bind_scalar("uniform", &bound, &BindOptions::new()).unwrap() + }; + let (zero, ten, twenty) = (seed_i32(0), seed_i32(10), seed_i32(20)); + let f = bound(&zero, &ten, &zero); + assert_eq!(f.return_type(), &DataType::Int32); + let out = f + .invoke_in(&[lit(&zero), lit(&ten), lit(&zero)], 1, 0) + .unwrap(); + assert_eq!(out.as_primitive::().values(), &[7]); + let f = bound(&ten, &twenty, &zero); + let out = f + .invoke_in(&[lit(&ten), lit(&twenty), lit(&zero)], 1, 0) + .unwrap(); + assert_eq!(out.as_primitive::().values(), &[17]); + // Partition 1 draws its own stream (`16`, then `15`). + let out = f + .invoke_in(&[lit(&ten), lit(&twenty), lit(&zero)], 2, 1) + .unwrap(); + assert_eq!(out.as_primitive::().values(), &[16, 15]); +} + +#[test] +fn randstr_matches_the_golden_strings() { + // `randstr(10, 0)` is `ceV0PXaR2I`, `randstr(3, 0)` is `ceV`. + let len = seed_i32(10); + let sd = seed_i32(0); + let bound = vec![ + BindArg::Literal(Arc::clone(&len)), + BindArg::Literal(Arc::clone(&sd)), + ]; + let f = super::bind_scalar("randstr", &bound, &BindOptions::new()).unwrap(); + assert_eq!(f.return_type(), &DataType::Utf8); + let out = f.invoke_in(&[lit(&len), lit(&sd)], 1, 0).unwrap(); + let texts: Vec> = out + .as_string::() + .iter() + .map(|s| s.map(str::to_owned)) + .collect(); + assert_eq!(texts, vec![Some("ceV0PXaR2I".to_owned())]); + // A fresh bind draws the same stream again. + let len3 = seed_i32(3); + let bound3 = vec![ + BindArg::Literal(Arc::clone(&len3)), + BindArg::Literal(Arc::clone(&sd)), + ]; + let g = super::bind_scalar("randstr", &bound3, &BindOptions::new()).unwrap(); + let out = g.invoke_in(&[lit(&len3), lit(&sd)], 1, 0).unwrap(); + let texts: Vec> = out + .as_string::() + .iter() + .map(|s| s.map(str::to_owned)) + .collect(); + assert_eq!(texts, vec![Some("ceV".to_owned())]); +} + +#[test] +fn randstr_rows_draw_consecutively() { + // One partition draws one stream: rows are consecutive slices of it + // (as a real Spark run with one partition shows). + let len = seed_i32(10); + let sd = seed_i32(0); + let bound = vec![ + BindArg::Literal(Arc::clone(&len)), + BindArg::Literal(Arc::clone(&sd)), + ]; + let f = super::bind_scalar("randstr", &bound, &BindOptions::new()).unwrap(); + let out = f.invoke_in(&[lit(&len), lit(&sd)], 3, 0).unwrap(); + let texts: Vec> = out + .as_string::() + .iter() + .map(|s| s.map(str::to_owned)) + .collect(); + assert_eq!( + texts, + vec![ + Some("ceV0PXaR2I".to_owned()), + Some("lBhHi56d0u".to_owned()), + Some("CtuwSOTLiL".to_owned()) + ] + ); + // The golden's second partition draws its own stream: with partitions + // `[0]`, `[1, 2]` the rows are `ceV0PXaR2I`, `fYxVfArnv7`, + // `iSIv0VT2XL`. + let out = f.invoke_in(&[lit(&len), lit(&sd)], 2, 1).unwrap(); + let texts: Vec> = out + .as_string::() + .iter() + .map(|s| s.map(str::to_owned)) + .collect(); + assert_eq!( + texts, + vec![Some("fYxVfArnv7".to_owned()), Some("iSIv0VT2XL".to_owned())] + ); +} + +#[test] +fn randstr_null_length_is_empty_and_negative_fails() { + use krishiv_common::error::ErrorClass; + let sd = seed_i32(0); + // A NULL length is the empty string. + let bound = vec![ + BindArg::Literal(null_i32()), + BindArg::Literal(Arc::clone(&sd)), + ]; + let f = super::bind_scalar("randstr", &bound, &BindOptions::new()).unwrap(); + let out = f.invoke_in(&[lit(&null_i32()), lit(&sd)], 2, 0).unwrap(); + assert_eq!(out.len(), 2); + assert!(out.as_string::().iter().all(|s| s == Some("")),); + // A negative length is `INVALID_PARAMETER_VALUE.LENGTH`. + let neg = seed_i32(-1); + let bound = vec![ + BindArg::Literal(Arc::clone(&neg)), + BindArg::Literal(Arc::clone(&sd)), + ]; + let f = super::bind_scalar("randstr", &bound, &BindOptions::new()).unwrap(); + let err = f.invoke_in(&[lit(&neg), lit(&sd)], 1, 0).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::InvalidParameterValueLength); + assert!(err.to_string().contains("`randstr`"), "{err}"); + assert!(err.to_string().contains("-1"), "{err}"); +} + +#[test] +fn uuids_have_shape_and_partitioned_streams() { + let sd = seed_i32(42); + let bound = vec![BindArg::Literal(Arc::clone(&sd))]; + let f = super::bind_scalar("uuid", &bound, &BindOptions::new()).unwrap(); + assert_eq!(f.return_type(), &DataType::Utf8); + let out = f.invoke_in(&[lit(&sd)], 2, 0).unwrap(); + let texts: Vec = out + .as_string::() + .iter() + .map(|s| s.unwrap().to_owned()) + .collect(); + assert_eq!(texts.len(), 2); + assert_ne!(texts[0], texts[1]); + for u in &texts { + assert_eq!(u.len(), 36, "{u}"); + assert_eq!(&u[8..9], "-"); + assert_eq!(&u[14..15], "4"); + assert!(matches!(&u[19..20], "8" | "9" | "a" | "b"), "{u}"); + } + // Another partition draws another stream. + let p1 = f.invoke_in(&[lit(&sd)], 2, 1).unwrap(); + let other: Vec = p1 + .as_string::() + .iter() + .map(|s| s.unwrap().to_owned()) + .collect(); + assert_ne!(texts, other); + // The unseeded form binds with no arguments. + let g = super::bind_scalar("uuid", &[], &BindOptions::new()).unwrap(); + let out = g.invoke_in(&[], 1, 0).unwrap(); + assert_eq!(out.len(), 1); +} + +#[test] +fn unseeded_draws_come_from_the_query_seed() { + // Two binds of `rand()` under one query seed agree; another query + // seed draws another stream. (Separate `BindOptions` hold separate + // salts; one query's binds draw different streams, below.) + let f = + super::bind_scalar("rand", &[], &BindOptions::new().with_query_start(1_000_000)).unwrap(); + let g = + super::bind_scalar("rand", &[], &BindOptions::new().with_query_start(1_000_000)).unwrap(); + assert_eq!( + f64s(&f.invoke_in(&[], 2, 0).unwrap()), + f64s(&g.invoke_in(&[], 2, 0).unwrap()) + ); + let other = BindOptions::new().with_query_start(2_000_000); + let h = super::bind_scalar("rand", &[], &other).unwrap(); + assert_ne!( + f64s(&f.invoke_in(&[], 2, 0).unwrap()), + f64s(&h.invoke_in(&[], 2, 0).unwrap()) + ); +} + +#[test] +fn unseeded_calls_in_one_query_draw_different_streams() { + // Two unseeded binds sharing one query's salt draw different streams + // (Spark's `ResolveRandomSeed` assigns each its own seed), and the + // same holds for `uuid` and `uniform` without a seed. + let opts = BindOptions::new().with_query_start(1_000_000); + let f = super::bind_scalar("rand", &[], &opts).unwrap(); + let g = super::bind_scalar("rand", &[], &opts).unwrap(); + assert_ne!( + f64s(&f.invoke_in(&[], 2, 0).unwrap()), + f64s(&g.invoke_in(&[], 2, 0).unwrap()) + ); + let u = super::bind_scalar("uuid", &[], &opts).unwrap(); + let v = super::bind_scalar("uuid", &[], &opts).unwrap(); + let texts = |f: &Arc| { + f.invoke_in(&[], 1, 0) + .unwrap() + .as_string::() + .iter() + .map(|s| s.unwrap().to_owned()) + .collect::>() + }; + assert_ne!(texts(&u), texts(&v)); +} + +#[test] +fn wrong_argument_types_do_not_bind() { + use krishiv_common::error::ErrorClass; + // A seed that is not an `int`/`bigint` is not bound (the analyzer + // reports `DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE` first): the error + // names the bad signature, not a missing function. + for t in [DataType::Float64, DataType::Utf8] { + let a = arrow_array::new_empty_array(&t); + let err = + super::bind_scalar("rand", &[BindArg::Column(t)], &BindOptions::new()).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::KrishivUnsupportedFeature); + assert!(err.to_string().contains("rand("), "{err}"); + let _ = a; + } + // `uniform` takes two or three arguments. + let err = super::bind_scalar("uniform", &[], &BindOptions::new()).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::KrishivUnsupportedFeature); + // `randstr` takes one or two. + let ty = || ints(&[Some(1)]).data_type().clone(); + let err = super::bind_scalar( + "randstr", + &[ty().into(), ty().into(), ty().into()], + &BindOptions::new(), + ) + .unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::KrishivUnsupportedFeature); + // A string length casts like Spark's implicit cast; a bad one fails + // the cast at runtime. + let len = || { + super::bind_scalar( + "randstr", + &[BindArg::Literal(strs(&[Some("10")]))], + &BindOptions::new(), + ) + .unwrap() + }; + let out = len() + .invoke_in(&[Arg::Scalar(strs(&[Some("10")]))], 1, 0) + .unwrap(); + assert_eq!(out.as_string::().value(0).len(), 10); + let err = len() + .invoke_in(&[Arg::Scalar(strs(&[Some("abc")]))], 1, 0) + .unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::CastInvalidInput); +} diff --git a/crates/krishiv-plan/src/analyzer/exprs.rs b/crates/krishiv-plan/src/analyzer/exprs.rs index 25335efce..ecddca509 100644 --- a/crates/krishiv-plan/src/analyzer/exprs.rs +++ b/crates/krishiv-plan/src/analyzer/exprs.rs @@ -876,6 +876,9 @@ impl Resolver<'_> { let (mut call_kids, ty) = self.typed_call(&entry, e, &kids[..n])?; let data_type = ty.data_type.clone().unwrap_or(DataType::Null); call_kids.extend(kids.get(n).cloned()); + // `typeof(x)` is Spark's `RuntimeReplaceable`; its type + // name is known once the call is typed. + let is_typeof = entry.id.name() == "typeof"; let call = ResolvedFunction { id: entry.id, kind: entry.kind, @@ -888,6 +891,21 @@ impl Resolver<'_> { }; let expr = rebuilt(ExprNode::Resolved(ResolvedExpr::Function(call))); let shown = self.call_shown(&f.name, &expr, kids, &call_kids); + // `typeof(x)` is Spark's `RuntimeReplaceable`: the analyzer + // replaces it with the type name literal, keeping the call + // as the shown form (the executor gets no type logic). + let (expr, shown) = if is_typeof { + let name = call_kids + .first() + .and_then(|k| k.ty.data_type.clone()) + .map_or_else(|| "void".to_owned(), |t| t.catalog_string()); + ( + rebuilt(ExprNode::Literal(Literal::String(name))), + Some(expr), + ) + } else { + (expr, shown) + }; Bound { agg: if entry.kind == FunctionKind::Aggregate && !window_fn { Some(e.origin) @@ -2480,6 +2498,16 @@ fn call( let d = if distinct { "DISTINCT " } else { "" }; format!("to_char({d}{})", k.join(", ")) } + // `current_database` and `current_schema` are one expression + // (Spark's `CurrentDatabase`, shown as `current_schema`); + // `current_catalog` shows without parentheses. + _ if name.eq_ignore_ascii_case("current_database") => { + let d = if distinct { "DISTINCT " } else { "" }; + format!("current_schema({d}{})", k.join(", ")) + } + _ if name.eq_ignore_ascii_case("current_catalog") && k.is_empty() && !distinct => { + "current_catalog".to_owned() + } // Single-argument math (`sec`, `csc`, `cot`, `exp`, ...) and // two-argument math (`power`, `atan2`, `hypot`) print upper case, // as Spark's `UnaryMathExpression` / `BinaryMathExpression` do; diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index 35008652b..5028bc4a2 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -1519,7 +1519,8 @@ fn unsupported_inputs_fail_explicitly() { old.profile = SparkProfile::V3_5; let randstr = LogicalPlan::table(["t"]).project(vec![Expr::function("randstr", vec![int(3)])]); assert_eq!(class(&run(randstr.clone(), &old)), "UNRESOLVED_ROUTINE"); - assert_eq!(class(&run(randstr, &ctx)), "KRISHIV_UNSUPPORTED_FEATURE"); + // `randstr` is implemented: it resolves on current versions. + assert_eq!(class(&run(randstr, &ctx)), "OK"); // The argument count is checked first, with Spark's `expectedNum` text. let upper = LogicalPlan::table(["t"]).project(vec![Expr::function("upper", vec![int(1), int(2)])]); @@ -1649,7 +1650,7 @@ fn function_entries_print_their_signature() { let uuid = ["uuid".to_owned()]; assert!(matches!( b.lookup(&uuid, SparkProfile::V4_2), - FunctionLookup::KnownUnsupported + FunctionLookup::Found(_) )); assert_eq!(b.version(), 0); } diff --git a/crates/krishiv-session/tests/misc_functions.rs b/crates/krishiv-session/tests/misc_functions.rs new file mode 100644 index 000000000..7dc367144 --- /dev/null +++ b/crates/krishiv-session/tests/misc_functions.rs @@ -0,0 +1,481 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The misc-function goldens end to end: every `misc-functions.sql`, +//! `random.sql` and `current_database_catalog.sql` statement that uses only +//! the implemented functions matches Spark 4.2.0 exactly (schema, rows, +//! error class and parameters), and the nondeterministic functions are +//! deterministic per partition. +//! +//! Two multi-row statements depend on Spark CI's partitioning (two slices +//! for three `VALUES` rows, then sorted output): the engine runs tiny scans +//! in one partition, so those exact row orders are asserted sorted where the +//! value multiset agrees, and at the operator level otherwise (see the +//! `#[test]` docs). +#![cfg(test)] + +use arrow_array::cast::AsArray; +use arrow_array::{Array, RecordBatch}; +use arrow_schema::DataType as ArrowType; +use arrow_schema::{Field, Schema}; +use krishiv_common::error::{ErrorClass, KrishivError}; +use krishiv_session::Session; +use std::sync::Arc; + +fn block_on(f: impl Future) -> T { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(f) +} + +/// The schema names and every row of `sql`, values as strings (`~` for +/// NULL), rows in output order. +fn rows(s: &Session, sql: &str) -> (Vec, Vec>) { + let df = block_on(s.sql(sql)).unwrap(); + let schema = block_on(df.schema()).unwrap(); + let names: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); + let batches = block_on(df.collect()).unwrap(); + let mut out = Vec::new(); + for b in &batches { + for i in 0..b.num_rows() { + let mut row = Vec::new(); + for c in b.columns() { + if c.is_null(i) { + row.push("~".to_owned()); + continue; + } + row.push(match c.data_type() { + ArrowType::Utf8View => c.as_string_view().value(i).to_owned(), + ArrowType::Utf8 => c.as_string::().value(i).to_owned(), + ArrowType::Int32 => c + .as_primitive::() + .value(i) + .to_string(), + ArrowType::Int64 => c + .as_primitive::() + .value(i) + .to_string(), + ArrowType::Float32 => c + .as_primitive::() + .value(i) + .to_string(), + ArrowType::Float64 => c + .as_primitive::() + .value(i) + .to_string(), + ArrowType::Boolean => c.as_boolean().value(i).to_string(), + ArrowType::Decimal128(_, s) => { + let v = c + .as_primitive::() + .value(i); + render_decimal(v, *s) + } + t => panic!("{sql}: unexpected type {t}"), + }); + } + out.push(row); + } + } + (names, out) +} + +fn session() -> Session { + Session::new().unwrap() +} + +/// Renders a decimal `(unscaled, scale)` as Spark does. +fn render_decimal(unscaled: i128, scale: i8) -> String { + let neg = unscaled < 0; + let digits = unscaled.unsigned_abs().to_string(); + let s = scale.cast_unsigned() as usize; + let text = if s == 0 { + digits + } else if digits.len() > s { + let (i, f) = digits.split_at(digits.len() - s); + format!("{i}.{f}") + } else { + format!("0.{:0>width$}", digits, width = s) + }; + format!("{}{text}", if neg { "-" } else { "" }) +} + +/// A session with a three-row table (`SELECT ... FROM vals` stands in for +/// the goldens' `FROM VALUES (0), (1), (2)`, which the engine does not +/// support yet). +fn values_session() -> Session { + use arrow_array::Int32Array; + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "col", + ArrowType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![0, 1, 2]))], + ) + .unwrap(); + s.register_batches("vals", schema, vec![batch]).unwrap(); + s +} + +fn error(s: &Session, sql: &str) -> KrishivError { + block_on(async { s.sql(sql).await?.collect().await }).unwrap_err() +} + +#[test] +fn typeof_names_the_analyzers_type() { + let s = session(); + // `select typeof(null)` in `misc-functions.sql`. + assert_eq!( + rows(&s, "select typeof(null)"), + ( + vec!["typeof(NULL)".to_owned()], + vec![vec!["void".to_owned()]] + ) + ); + assert_eq!( + rows(&s, "select typeof(true)"), + ( + vec!["typeof(true)".to_owned()], + vec![vec!["boolean".to_owned()]] + ) + ); + assert_eq!( + rows(&s, "select typeof(1Y), typeof(1S), typeof(1), typeof(1L)"), + ( + vec![ + "typeof(1)".to_owned(), + "typeof(1)".to_owned(), + "typeof(1)".to_owned(), + "typeof(1)".to_owned(), + ], + vec![vec![ + "tinyint".to_owned(), + "smallint".to_owned(), + "int".to_owned(), + "bigint".to_owned(), + ]] + ) + ); + assert_eq!( + rows( + &s, + "select typeof(cast(1.0 as float)), typeof(1.0D), typeof(1.2)" + ), + ( + vec![ + "typeof(CAST(1.0 AS FLOAT))".to_owned(), + "typeof(1.0)".to_owned(), + "typeof(1.2)".to_owned(), + ], + vec![vec![ + "float".to_owned(), + "double".to_owned(), + "decimal(2,1)".to_owned(), + ]] + ) + ); + assert_eq!( + rows(&s, "select typeof(x'ABCD'), typeof('SPARK')"), + ( + vec!["typeof(X'ABCD')".to_owned(), "typeof(SPARK)".to_owned()], + vec![vec!["binary".to_owned(), "string".to_owned()]] + ) + ); + assert_eq!( + rows( + &s, + "select typeof(date '1986-05-23'), typeof(timestamp '1986-05-23')" + ), + ( + vec![ + "typeof(DATE '1986-05-23')".to_owned(), + "typeof(TIMESTAMP '1986-05-23 00:00:00')".to_owned(), + ], + vec![vec!["date".to_owned(), "timestamp".to_owned()]] + ) + ); +} + +#[test] +fn typeof_needs_no_execution() { + // The rewrite is a literal: the argument is never evaluated, so a + // failing argument still names its type (with the analyzer's casts, as + // Spark's `toPrettySQL` shows them). + let s = session(); + assert_eq!( + rows(&s, "select typeof(1 / 0)").0, + vec!["typeof((CAST(1 AS DOUBLE) / CAST(0 AS DOUBLE)))".to_owned()] + ); +} + +#[test] +fn current_database_and_catalog() { + // `current_database_catalog.sql`, exactly. + let s = session(); + assert_eq!( + rows( + &s, + "select current_database(), current_schema(), current_catalog()" + ), + ( + vec![ + "current_schema()".to_owned(), + "current_schema()".to_owned(), + "current_catalog".to_owned(), + ], + vec![vec![ + "default".to_owned(), + "default".to_owned(), + "spark_catalog".to_owned(), + ]] + ) + ); + assert_eq!( + rows(&s, "select version()").1, + vec![vec!["4.2.0".to_owned()]] + ); +} + +#[test] +fn rand_and_randn_match_spark() { + let s = session(); + // `random.sql`, exactly. + assert_eq!( + rows(&s, "SELECT rand(0)").1, + vec![vec!["0.7604953758285915".to_owned()]] + ); + assert_eq!( + rows(&s, "SELECT rand(cast(3 / 7 AS int))").1, + vec![vec!["0.7604953758285915".to_owned()]] + ); + assert_eq!( + rows(&s, "SELECT rand(NULL)").1, + vec![vec!["0.7604953758285915".to_owned()]] + ); + assert_eq!( + rows(&s, "SELECT rand(cast(NULL AS int))").1, + vec![vec!["0.7604953758285915".to_owned()]] + ); + assert_eq!( + rows(&s, "SELECT randn(0L)").1, + vec![vec!["1.6034991609278433".to_owned()]] + ); + assert_eq!( + rows(&s, "SELECT randn(cast(3 / 7 AS long))").1, + vec![vec!["1.6034991609278433".to_owned()]] + ); + assert_eq!( + rows(&s, "SELECT randn(NULL)").1, + vec![vec!["1.6034991609278433".to_owned()]] + ); + assert_eq!( + rows(&s, "SELECT randn(cast(NULL AS long))").1, + vec![vec!["1.6034991609278433".to_owned()]] + ); + // The seed types Spark rejects. + let e = error(&s, "SELECT rand(1.0)"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(e.to_string().contains("(\"INT\" or \"BIGINT\")"), "{e}"); + let e = error(&s, "SELECT rand('1')"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(e.to_string().contains("\"rand(1)\""), "{e}"); + // A non-foldable seed is refused. + let v = values_session(); + let e = error(&v, "SELECT rand(col) FROM vals"); + assert_eq!(e.error_class(), ErrorClass::SeedExpressionIsUnfoldable); + let e = error(&s, "SELECT rand(id) FROM (SELECT 1 AS id)"); + assert_eq!(e.error_class(), ErrorClass::SeedExpressionIsUnfoldable); +} + +#[test] +fn uniform_matches_spark() { + let s = session(); + let value = |sql: &str| rows(&s, sql).1.into_iter().flatten().collect::>(); + assert_eq!(value("SELECT uniform(0, 1, 0) AS result"), ["0"]); + assert_eq!(value("SELECT uniform(0, 10, 0) AS result"), ["7"]); + assert_eq!(value("SELECT uniform(0L, 10L, 0) AS result"), ["7"]); + assert_eq!(value("SELECT uniform(10, 20, 0) AS result"), ["17"]); + assert_eq!( + value("SELECT uniform(10.0F, 20.0F, 0) AS result"), + ["17.604954"] + ); + assert_eq!( + value( + "SELECT uniform(cast(10 as decimal(10, 3)), cast(20 as decimal(10, 3)), 0) AS result" + ), + ["17.605"] + ); + assert_eq!( + value( + "SELECT uniform(cast(10 as decimal(10, 3)), cast(20 as decimal(11, 4)), 0) AS result" + ), + ["17.6050"] + ); + assert_eq!( + value("SELECT uniform(10.0D, 20.0D, CAST(3 / 7 AS LONG)) AS result"), + ["17.604953758285916"] + ); + assert_eq!(value("SELECT uniform(-10L, 10L, 0) AS result"), ["5"]); + assert_eq!(value("SELECT uniform(-20L, -10L, -10) AS result"), ["-17"]); + // NULL bounds are NULL results; a NULL seed reads as 0. + assert_eq!(value("SELECT uniform(NULL, 1, 0) AS result"), ["~"]); + assert_eq!( + value("SELECT uniform(cast(NULL AS int), 1, 0) AS result"), + ["~"] + ); + assert_eq!(value("SELECT uniform(0, NULL, 0) AS result"), ["~"]); + assert_eq!(value("SELECT uniform(0, 1, NULL) AS result"), ["0"]); + assert_eq!( + value("SELECT uniform(0, 1, cast(NULL as int)) AS result"), + ["0"] + ); + // The seed and bound checks Spark reports. + let v = values_session(); + let e = error(&v, "SELECT uniform(10, 20, col) AS result FROM vals"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchNonFoldableInput + ); + assert!(e.to_string().contains("`seed`"), "{e}"); + let e = error(&s, "SELECT uniform(10) AS result"); + assert_eq!(e.error_class(), ErrorClass::WrongNumArgsWithoutSuggestion); + let e = error(&s, "SELECT uniform(10.0F, 20.0F, 0.0F) AS result"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(e.to_string().contains("third"), "{e}"); + let e = error(&s, "SELECT uniform('abc', 10, 0) AS result"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(e.to_string().contains("\"NUMERIC\""), "{e}"); + // Several partitions draw several streams: the engine runs the tiny + // scan in one partition, so the rows are one stream's first draws + // (Spark CI splits the rows over two partitions and sorts). + let v = values_session(); + let mut got = rows(&v, "SELECT uniform(10, 20, 0) AS result FROM vals") + .1 + .into_iter() + .flatten() + .collect::>(); + got.sort(); + assert_eq!(got, ["10", "15", "17"]); +} + +#[test] +fn randstr_matches_spark() { + let s = session(); + let value = |sql: &str| rows(&s, sql).1.into_iter().flatten().collect::>(); + assert_eq!(value("SELECT randstr(1, 0) AS result"), ["c"]); + assert_eq!(value("SELECT randstr(5, 0) AS result"), ["ceV0P"]); + assert_eq!(value("SELECT randstr(10, 0) AS result"), ["ceV0PXaR2I"]); + assert_eq!(value("SELECT randstr(10S, 0) AS result"), ["ceV0PXaR2I"]); + assert_eq!( + value("SELECT randstr(CAST(10 AS TINYINT), 0) AS result"), + ["ceV0PXaR2I"] + ); + assert_eq!( + value("SELECT randstr(CAST(10 AS BIGINT), 0) AS result"), + ["ceV0PXaR2I"] + ); + assert_eq!(value("SELECT randstr(1.0F, 0) AS result"), ["c"]); + assert_eq!(value("SELECT randstr(1.0D, 0) AS result"), ["c"]); + assert_eq!( + value("SELECT randstr(cast(1 AS DECIMAL(10, 2)), 0) AS result"), + ["c"] + ); + assert_eq!(value("SELECT randstr(1, -1) AS result"), ["S"]); + assert_eq!(value("SELECT randstr(NULL, 0) AS result"), [""]); + assert_eq!(value("SELECT randstr(0, NULL) AS result"), [""]); + let v = values_session(); + let e = error(&v, "SELECT randstr(col, 0) AS result FROM vals"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchNonFoldableInput + ); + let e = error(&s, "SELECT randstr(10, 0, 1) AS result"); + assert_eq!(e.error_class(), ErrorClass::WrongNumArgsWithoutSuggestion); + let e = error(&s, "SELECT randstr(-1, 0) AS result"); + assert_eq!(e.error_class(), ErrorClass::InvalidParameterValueLength); + // The engine's tiny scan is one partition: consecutive draws of one + // stream (Spark CI splits the rows over two partitions and sorts). + let v = values_session(); + assert_eq!( + rows(&v, "SELECT randstr(10, 0) AS result FROM vals").1, + [ + vec!["ceV0PXaR2I".to_owned()], + vec!["lBhHi56d0u".to_owned()], + vec!["CtuwSOTLiL".to_owned()] + ] + ); +} + +#[test] +fn nondeterministic_queries_are_deterministic_per_partition() { + let s = session(); + let value = |sql: &str| rows(&s, sql).1.into_iter().flatten().collect::>(); + // The same query twice draws the same stream (one partition). + let once = value("SELECT rand(0), randn(0), uuid(42), monotonically_increasing_id()"); + let twice = value("SELECT rand(0), randn(0), uuid(42), monotonically_increasing_id()"); + assert_eq!(once, twice); + // An unseeded call is not NULL. + assert_eq!(value("SELECT randstr(10) IS NOT NULL AS result"), ["true"]); + assert_eq!( + value("SELECT uniform(10, 20.0F) IS NOT NULL AS result"), + ["true"] + ); + // `uuid()` draws from the query seed: two queries differ. + let a = value("SELECT uuid()"); + let b = value("SELECT uuid()"); + assert_ne!(a, b); + // Two unseeded calls in one query draw different streams. + let both = rows(&s, "SELECT rand() AS a, rand() AS b").1; + assert_ne!(both[0][0], both[0][1]); + // Partition functions in one partition. + assert_eq!(value("SELECT spark_partition_id()"), ["0"]); + assert_eq!(value("SELECT monotonically_increasing_id()"), ["0"]); + assert_eq!(value("SELECT input_file_name()"), [""]); + assert_eq!(value("SELECT input_file_block_start()"), ["-1"]); + assert_eq!(value("SELECT input_file_block_length()"), ["-1"]); + // A stateful predicate sees the same partition stream. + let v = values_session(); + assert_eq!( + rows(&v, "SELECT col FROM vals WHERE rand(0) > 0.5").1, + [vec!["0".to_owned()], vec!["1".to_owned()]] + ); +} + +#[test] +fn misc_rejects_bad_calls() { + let s = session(); + // `uuid` wants a plain literal seed. + let e = error(&s, "SELECT uuid(1 + 1)"); + assert_eq!(e.error_class(), ErrorClass::InvalidParameterValueLong); + // `typeof` takes exactly one argument. + let e = error(&s, "SELECT typeof()"); + assert_eq!(e.error_class(), ErrorClass::WrongNumArgsWithoutSuggestion); + for sql in [ + "SELECT version(1)", + "SELECT current_database(1)", + "SELECT monotonically_increasing_id(1)", + "SELECT input_file_name('x')", + "SELECT spark_partition_id(1)", + ] { + let e = error(&s, sql); + assert_eq!( + e.error_class(), + ErrorClass::WrongNumArgsWithoutSuggestion, + "{sql}" + ); + } +} diff --git a/tests/fixtures/analysis/registry-matrix-manifest.json b/tests/fixtures/analysis/registry-matrix-manifest.json index 058996fa6..b937ae267 100644 --- a/tests/fixtures/analysis/registry-matrix-manifest.json +++ b/tests/fixtures/analysis/registry-matrix-manifest.json @@ -1,5 +1,5 @@ { - "cases": 12566, + "cases": 13040, "entries": [ { "id": "c1a16bec609/ansi_off", @@ -8089,6 +8089,206 @@ "id": "raise_error#fcc723f51e83/ansi_on", "reason": "SELECT raise_error(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, + { + "id": "rand#114bc082913c/ansi_off", + "reason": "SELECT rand(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#114bc082913c/ansi_on", + "reason": "SELECT rand(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#16b39c8effd3/ansi_off", + "reason": "SELECT rand(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#16b39c8effd3/ansi_on", + "reason": "SELECT rand(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#2883ca45911e/ansi_off", + "reason": "SELECT rand(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#2883ca45911e/ansi_on", + "reason": "SELECT rand(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#5dd406b51bff/ansi_off", + "reason": "SELECT rand(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#5dd406b51bff/ansi_on", + "reason": "SELECT rand(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#a4f856765524/ansi_off", + "reason": "SELECT rand(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "rand#a4f856765524/ansi_on", + "reason": "SELECT rand(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#28f8edef6b0c/ansi_off", + "reason": "SELECT randn(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#28f8edef6b0c/ansi_on", + "reason": "SELECT randn(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#8df4dfcb761d/ansi_off", + "reason": "SELECT randn(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#8df4dfcb761d/ansi_on", + "reason": "SELECT randn(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#905f1cb314d6/ansi_off", + "reason": "SELECT randn(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#905f1cb314d6/ansi_on", + "reason": "SELECT randn(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#9df340e4c4ed/ansi_off", + "reason": "SELECT randn(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#9df340e4c4ed/ansi_on", + "reason": "SELECT randn(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#db9474536e6f/ansi_off", + "reason": "SELECT randn(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randn#db9474536e6f/ansi_on", + "reason": "SELECT randn(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#56f3c76a60d9/ansi_off", + "reason": "SELECT random(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#56f3c76a60d9/ansi_on", + "reason": "SELECT random(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#6a872774b235/ansi_off", + "reason": "SELECT random(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#6a872774b235/ansi_on", + "reason": "SELECT random(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#73bb8aca50b7/ansi_off", + "reason": "SELECT random(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#73bb8aca50b7/ansi_on", + "reason": "SELECT random(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#aafd575271fd/ansi_off", + "reason": "SELECT random(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#aafd575271fd/ansi_on", + "reason": "SELECT random(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#b218d468723d/ansi_off", + "reason": "SELECT random(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "random#b218d468723d/ansi_on", + "reason": "SELECT random(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#175c23e922dc/ansi_off", + "reason": "SELECT randstr(3, INTERVAL '1-2' YEAR TO MONTH) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#175c23e922dc/ansi_on", + "reason": "SELECT randstr(3, INTERVAL '1-2' YEAR TO MONTH) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#380822c97372/ansi_off", + "reason": "SELECT randstr(named_struct('a', 1, 'b', 'x'), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#380822c97372/ansi_on", + "reason": "SELECT randstr(named_struct('a', 1, 'b', 'x'), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#42f50df56ea3/ansi_off", + "reason": "SELECT randstr(3, named_struct('a', 1, 'b', 'x')) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#42f50df56ea3/ansi_on", + "reason": "SELECT randstr(3, named_struct('a', 1, 'b', 'x')) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#55df7b2e1334/ansi_off", + "reason": "SELECT randstr(3, INTERVAL '1 02:03:04' DAY TO SECOND) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#55df7b2e1334/ansi_on", + "reason": "SELECT randstr(3, INTERVAL '1 02:03:04' DAY TO SECOND) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#72a2793557e6/ansi_off", + "reason": "SELECT randstr(array(1, 2), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#72a2793557e6/ansi_on", + "reason": "SELECT randstr(array(1, 2), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#78a2b0d4ee1a/ansi_off", + "reason": "SELECT randstr(3, array(1, 2)) AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#78a2b0d4ee1a/ansi_on", + "reason": "SELECT randstr(3, array(1, 2)) AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#9e4c25337d0d/ansi_off", + "reason": "SELECT randstr(map('a', 1), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#9e4c25337d0d/ansi_on", + "reason": "SELECT randstr(map('a', 1), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#a0e222978047/ansi_off", + "reason": "SELECT randstr(INTERVAL '1 02:03:04' DAY TO SECOND, 0) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#a0e222978047/ansi_on", + "reason": "SELECT randstr(INTERVAL '1 02:03:04' DAY TO SECOND, 0) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#d06ac69e5e3a/ansi_off", + "reason": "SELECT randstr(3, map('a', 1)) AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#d06ac69e5e3a/ansi_on", + "reason": "SELECT randstr(3, map('a', 1)) AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#e2c6f88c4702/ansi_off", + "reason": "SELECT randstr(INTERVAL '1-2' YEAR TO MONTH, 0) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "randstr#e2c6f88c4702/ansi_on", + "reason": "SELECT randstr(INTERVAL '1-2' YEAR TO MONTH, 0) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, { "id": "rank#a1a83b14b6f8/ansi_off", "reason": "SELECT a, b, rank(b) OVER (PARTITION BY a ORDER BY b) FROM VALUES ('A1', 2), ('A1', 1), ('A2', 3), ('A1', 1) tab(a, b) | parse: KRISHIV_UNSUPPORTED_FEATURE: VALUES" @@ -11817,6 +12017,94 @@ "id": "try_to_number#c4e8862c7284/ansi_on", "reason": "SELECT try_to_number('454', TIMESTAMP'2024-02-29 12:34:56.123456') | DATATYPE_MISMATCH.NON_FOLDABLE_INPUT where Spark has INVALID_FORMAT.UNEXPECTED_TOKEN" }, + { + "id": "typeof#737efe3cca2b/ansi_off", + "reason": "SELECT typeof(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map" + }, + { + "id": "typeof#737efe3cca2b/ansi_on", + "reason": "SELECT typeof(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map" + }, + { + "id": "typeof#7b81ac5ac6dd/ansi_off", + "reason": "SELECT typeof(CAST(NULL AS array)) | parse: KRISHIV_UNSUPPORTED_FEATURE: nested data types (\"ARRAY\")" + }, + { + "id": "typeof#7b81ac5ac6dd/ansi_on", + "reason": "SELECT typeof(CAST(NULL AS array)) | parse: KRISHIV_UNSUPPORTED_FEATURE: nested data types (\"ARRAY\")" + }, + { + "id": "typeof#85119c3e6111/ansi_off", + "reason": "SELECT typeof(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals" + }, + { + "id": "typeof#85119c3e6111/ansi_on", + "reason": "SELECT typeof(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals" + }, + { + "id": "typeof#86964662cdab/ansi_off", + "reason": "SELECT typeof(array(1)) | KRISHIV_UNSUPPORTED_FEATURE: function array" + }, + { + "id": "typeof#86964662cdab/ansi_on", + "reason": "SELECT typeof(array(1)) | KRISHIV_UNSUPPORTED_FEATURE: function array" + }, + { + "id": "typeof#8a4ed8746829/ansi_off", + "reason": "SELECT typeof(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct" + }, + { + "id": "typeof#8a4ed8746829/ansi_on", + "reason": "SELECT typeof(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct" + }, + { + "id": "typeof#8b64929c858b/ansi_off", + "reason": "SELECT typeof(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array" + }, + { + "id": "typeof#8b64929c858b/ansi_on", + "reason": "SELECT typeof(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array" + }, + { + "id": "typeof#94be9d2b4025/ansi_off", + "reason": "SELECT typeof(CAST(array() AS array)) | parse: KRISHIV_UNSUPPORTED_FEATURE: nested data types (\"ARRAY\")" + }, + { + "id": "typeof#94be9d2b4025/ansi_on", + "reason": "SELECT typeof(CAST(array() AS array)) | parse: KRISHIV_UNSUPPORTED_FEATURE: nested data types (\"ARRAY\")" + }, + { + "id": "typeof#bea7c4b2a1ca/ansi_off", + "reason": "SELECT typeof(CAST(array(NULL) AS array)) | parse: KRISHIV_UNSUPPORTED_FEATURE: nested data types (\"ARRAY\")" + }, + { + "id": "typeof#bea7c4b2a1ca/ansi_on", + "reason": "SELECT typeof(CAST(array(NULL) AS array)) | parse: KRISHIV_UNSUPPORTED_FEATURE: nested data types (\"ARRAY\")" + }, + { + "id": "typeof#e22394e72946/ansi_off", + "reason": "SELECT typeof(CAST(NULL AS int)) | string where Spark has string not null" + }, + { + "id": "typeof#e22394e72946/ansi_on", + "reason": "SELECT typeof(CAST(NULL AS int)) | string where Spark has string not null" + }, + { + "id": "typeof#e571d1babdbc/ansi_off", + "reason": "SELECT typeof(NULL) | string where Spark has string not null" + }, + { + "id": "typeof#e571d1babdbc/ansi_on", + "reason": "SELECT typeof(NULL) | string where Spark has string not null" + }, + { + "id": "typeof#ef1847dcc457/ansi_off", + "reason": "SELECT typeof(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals" + }, + { + "id": "typeof#ef1847dcc457/ansi_on", + "reason": "SELECT typeof(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals" + }, { "id": "ucase#05ac574f466d/ansi_off", "reason": "SELECT ucase(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals" @@ -12073,6 +12361,190 @@ "id": "unhex#ffcacbfe03e1/ansi_on", "reason": "SELECT decode(unhex(CAST(X'C328FF' AS STRING)), 'UTF-8') | KRISHIV_UNSUPPORTED_FEATURE: function decode" }, + { + "id": "uniform#183cfc5f631a/ansi_off", + "reason": "SELECT uniform(NULL, 20, 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#183cfc5f631a/ansi_on", + "reason": "SELECT uniform(NULL, 20, 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#1e99715e4fd3/ansi_off", + "reason": "SELECT uniform(10, INTERVAL '1 02:03:04' DAY TO SECOND, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#1e99715e4fd3/ansi_on", + "reason": "SELECT uniform(10, INTERVAL '1 02:03:04' DAY TO SECOND, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#30c2e2b398e9/ansi_off", + "reason": "SELECT uniform(array(1, 2), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#30c2e2b398e9/ansi_on", + "reason": "SELECT uniform(array(1, 2), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#310b01a5b4ae/ansi_off", + "reason": "SELECT uniform(10, INTERVAL '1-2' YEAR TO MONTH, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#310b01a5b4ae/ansi_on", + "reason": "SELECT uniform(10, INTERVAL '1-2' YEAR TO MONTH, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#3fd714d3fbcf/ansi_off", + "reason": "SELECT uniform(10, map('a', 1), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#3fd714d3fbcf/ansi_on", + "reason": "SELECT uniform(10, map('a', 1), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#454025458e8a/ansi_off", + "reason": "SELECT uniform(10, 20, NULL) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#454025458e8a/ansi_on", + "reason": "SELECT uniform(10, 20, NULL) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#523b3fcbd4d8/ansi_off", + "reason": "SELECT uniform(10, 20, named_struct('a', 1, 'b', 'x')) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#523b3fcbd4d8/ansi_on", + "reason": "SELECT uniform(10, 20, named_struct('a', 1, 'b', 'x')) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#5656d353ced6/ansi_off", + "reason": "SELECT uniform(INTERVAL '1-2' YEAR TO MONTH, 20, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#5656d353ced6/ansi_on", + "reason": "SELECT uniform(INTERVAL '1-2' YEAR TO MONTH, 20, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#5ff5470f6623/ansi_off", + "reason": "SELECT uniform(10, 20, array(1, 2)) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#5ff5470f6623/ansi_on", + "reason": "SELECT uniform(10, 20, array(1, 2)) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#681f70cb51f1/ansi_off", + "reason": "SELECT uniform(map('a', 1), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#681f70cb51f1/ansi_on", + "reason": "SELECT uniform(map('a', 1), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#6afd046a5774/ansi_off", + "reason": "SELECT uniform(10, 20, INTERVAL '1-2' YEAR TO MONTH) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#6afd046a5774/ansi_on", + "reason": "SELECT uniform(10, 20, INTERVAL '1-2' YEAR TO MONTH) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#8f340576d087/ansi_off", + "reason": "SELECT uniform(10, named_struct('a', 1, 'b', 'x'), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#8f340576d087/ansi_on", + "reason": "SELECT uniform(10, named_struct('a', 1, 'b', 'x'), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#91babb0d6174/ansi_off", + "reason": "SELECT uniform(named_struct('a', 1, 'b', 'x'), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#91babb0d6174/ansi_on", + "reason": "SELECT uniform(named_struct('a', 1, 'b', 'x'), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#96fbf5c2c8f2/ansi_off", + "reason": "SELECT uniform(10, CAST(1.5 AS DECIMAL(38,18)), 0) > 0 AS result | boolean not null where Spark has boolean" + }, + { + "id": "uniform#96fbf5c2c8f2/ansi_on", + "reason": "SELECT uniform(10, CAST(1.5 AS DECIMAL(38,18)), 0) > 0 AS result | boolean not null where Spark has boolean" + }, + { + "id": "uniform#a51977da5bbb/ansi_off", + "reason": "SELECT uniform(10, 20, CAST(NULL AS int)) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#a51977da5bbb/ansi_on", + "reason": "SELECT uniform(10, 20, CAST(NULL AS int)) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#a9c82a62607d/ansi_off", + "reason": "SELECT uniform(10, CAST(NULL AS int), 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#a9c82a62607d/ansi_on", + "reason": "SELECT uniform(10, CAST(NULL AS int), 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#ab6b99e0333b/ansi_off", + "reason": "SELECT uniform(10, NULL, 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#ab6b99e0333b/ansi_on", + "reason": "SELECT uniform(10, NULL, 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#ac2a6eb487eb/ansi_off", + "reason": "SELECT uniform(10, array(1, 2), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#ac2a6eb487eb/ansi_on", + "reason": "SELECT uniform(10, array(1, 2), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function array where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#d54a4a610339/ansi_off", + "reason": "SELECT uniform(10, 20, map('a', 1)) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#d54a4a610339/ansi_on", + "reason": "SELECT uniform(10, 20, map('a', 1)) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#d942efa11da8/ansi_off", + "reason": "SELECT uniform(CAST(1.5 AS DECIMAL(38,18)), 20, 0) > 0 AS result | boolean not null where Spark has boolean" + }, + { + "id": "uniform#d942efa11da8/ansi_on", + "reason": "SELECT uniform(CAST(1.5 AS DECIMAL(38,18)), 20, 0) > 0 AS result | boolean not null where Spark has boolean" + }, + { + "id": "uniform#e6040ed43075/ansi_off", + "reason": "SELECT uniform(INTERVAL '1 02:03:04' DAY TO SECOND, 20, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#e6040ed43075/ansi_on", + "reason": "SELECT uniform(INTERVAL '1 02:03:04' DAY TO SECOND, 20, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#f476aa44bc4b/ansi_off", + "reason": "SELECT uniform(CAST(NULL AS int), 20, 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#f476aa44bc4b/ansi_on", + "reason": "SELECT uniform(CAST(NULL AS int), 20, 0) > 0 AS result | boolean where Spark has boolean not null" + }, + { + "id": "uniform#fa2e75b8ab98/ansi_off", + "reason": "SELECT uniform(10, 20, INTERVAL '1 02:03:04' DAY TO SECOND) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, + { + "id": "uniform#fa2e75b8ab98/ansi_on", + "reason": "SELECT uniform(10, 20, INTERVAL '1 02:03:04' DAY TO SECOND) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + }, { "id": "upper#04b381e12869/ansi_off", "reason": "SELECT upper(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" @@ -12968,11 +13440,26 @@ "function": "curdate", "pass": 2 }, + { + "cases": 2, + "function": "current_catalog", + "pass": 2 + }, + { + "cases": 2, + "function": "current_database", + "pass": 2 + }, { "cases": 4, "function": "current_date", "pass": 4 }, + { + "cases": 2, + "function": "current_schema", + "pass": 2 + }, { "cases": 4, "function": "current_timestamp", @@ -13118,6 +13605,21 @@ "function": "ifnull", "pass": 32 }, + { + "cases": 2, + "function": "input_file_block_length", + "pass": 2 + }, + { + "cases": 2, + "function": "input_file_block_start", + "pass": 2 + }, + { + "cases": 2, + "function": "input_file_name", + "pass": 2 + }, { "cases": 54, "function": "isnan", @@ -13238,6 +13740,11 @@ "function": "mod", "pass": 76 }, + { + "cases": 2, + "function": "monotonically_increasing_id", + "pass": 2 + }, { "cases": 54, "function": "month", @@ -13328,6 +13835,26 @@ "function": "raise_error", "pass": 44 }, + { + "cases": 54, + "function": "rand", + "pass": 44 + }, + { + "cases": 54, + "function": "randn", + "pass": 44 + }, + { + "cases": 54, + "function": "random", + "pass": 44 + }, + { + "cases": 90, + "function": "randstr", + "pass": 70 + }, { "cases": 2, "function": "rank", @@ -13458,6 +13985,11 @@ "function": "sinh", "pass": 40 }, + { + "cases": 2, + "function": "spark_partition_id", + "pass": 2 + }, { "cases": 50, "function": "sqrt", @@ -13543,6 +14075,11 @@ "function": "try_to_number", "pass": 90 }, + { + "cases": 56, + "function": "typeof", + "pass": 34 + }, { "cases": 54, "function": "ucase", @@ -13553,6 +14090,11 @@ "function": "unhex", "pass": 0 }, + { + "cases": 146, + "function": "uniform", + "pass": 100 + }, { "cases": 54, "function": "upper", @@ -13563,6 +14105,16 @@ "function": "user", "pass": 2 }, + { + "cases": 2, + "function": "uuid", + "pass": 2 + }, + { + "cases": 2, + "function": "version", + "pass": 2 + }, { "cases": 150, "function": "xxhash64", @@ -13589,5 +14141,5 @@ "pass": 40 } ], - "pass": 9370 + "pass": 9726 } From 088ceafbed61a98160025ef5af2b3d4725bf21dd Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Sun, 25 Oct 2026 02:27:25 +0100 Subject: [PATCH 2/7] Address review findings on misc random functions Fix per-partition stream seeding: record the base seed once and derive slot j as base + j, so a driver binding a higher partition first no longer seeds lower slots with its own seed. Callers pass the unfolded base seed. Fail closed on operator state mismatch: Filter, Project and Expand now return an internal error on foreign local state instead of silently using partition 0; Expand carries the driver partition to its inner projections. Share the uniform result-type rule between the analyzer and the executor binding through one kind-based decision function, with a cross-layer consistency test over every pair of bound kinds. Clarify golden coverage: session tests assert the engine single partition draws and are renamed accordingly; the per-partition decomposition stays pinned at binding level. Add the missing golden value and error cases. Extend the unverified-values note to version/input_file_*, and document the partition-index bounds and the untracked per-partition state. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/operators/expand.rs | 36 ++++-- crates/krishiv-exec/src/operators/filter.rs | 7 +- crates/krishiv-exec/src/operators/project.rs | 7 +- crates/krishiv-exec/src/operators/tests.rs | 66 ++++++++--- crates/krishiv-functions/src/lib.rs | 4 + crates/krishiv-functions/src/registry/misc.rs | 79 ++++++------- .../krishiv-functions/src/registry/tests.rs | 50 ++++++++ crates/krishiv-functions/src/scalar/misc.rs | 11 +- crates/krishiv-functions/src/scalar/random.rs | 109 ++++++++---------- .../src/scalar/tests/random.rs | 15 +++ crates/krishiv-functions/src/uniform.rs | 83 +++++++++++++ .../krishiv-session/tests/misc_functions.rs | 56 +++++++-- 12 files changed, 370 insertions(+), 153 deletions(-) create mode 100644 crates/krishiv-functions/src/uniform.rs diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs index aedb7a718..fca4cae65 100644 --- a/crates/krishiv-exec/src/operators/expand.rs +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -8,7 +8,7 @@ use arrow_array::{ArrayRef, RecordBatch, new_null_array}; use arrow_schema::{DataType, SchemaRef}; use krishiv_types::batch::Batch; -use super::{Project, ScalarExpr, local_mut, no_state}; +use super::{PartitionState, Project, ScalarExpr, local_mut}; use crate::context::TaskContext; use crate::error::{Result, internal}; use crate::operator::{LocalState, Operator, OperatorOutput}; @@ -56,8 +56,14 @@ impl Expand { } } -/// The next projection to emit for the current input batch. -struct Cursor(usize); +/// The next projection to emit for the current input batch, and the +/// driver's partition: the projections evaluate row expressions, which draw +/// the partition's stream. +#[derive(Debug, Clone, Copy)] +struct State { + next: usize, + partition: u32, +} impl Operator for Expand { fn name(&self) -> &'static str { @@ -68,29 +74,35 @@ impl Operator for Expand { Arc::clone(&self.schema) } - fn local_state(&self, _ctx: &TaskContext) -> Result { - Ok(Box::new(Cursor(0))) + fn local_state(&self, ctx: &TaskContext) -> Result { + Ok(Box::new(State { + next: 0, + partition: ctx.partition(), + })) } fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { - let cursor = local_mut::(local, "Expand")?; + let state = local_mut::(local, "Expand")?; if input.num_rows() == 0 { - cursor.0 = 0; + state.next = 0; return Ok(OperatorOutput::NeedMoreInput(None)); } let project = self .projects - .get(cursor.0) + .get(state.next) .ok_or_else(|| internal("an expansion past its last projection".to_owned()))?; - let OperatorOutput::NeedMoreInput(Some(out)) = project.execute(input, &mut no_state())? + let mut partition: LocalState = Box::new(PartitionState { + partition: state.partition, + }); + let OperatorOutput::NeedMoreInput(Some(out)) = project.execute(input, &mut partition)? else { return Err(internal("a projection did not output a batch".to_owned())); }; - cursor.0 += 1; - if cursor.0 < self.projects.len() { + state.next += 1; + if state.next < self.projects.len() { return Ok(OperatorOutput::HaveMoreOutput(out)); } - cursor.0 = 0; + state.next = 0; Ok(OperatorOutput::NeedMoreInput(Some(out))) } } diff --git a/crates/krishiv-exec/src/operators/filter.rs b/crates/krishiv-exec/src/operators/filter.rs index c65ac360d..17a039cbd 100644 --- a/crates/krishiv-exec/src/operators/filter.rs +++ b/crates/krishiv-exec/src/operators/filter.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use arrow_schema::SchemaRef; use krishiv_types::batch::Batch; -use super::PartitionState; use super::expr::Predicate; +use super::{PartitionState, local_mut}; use crate::context::TaskContext; use crate::error::Result; use crate::operator::{LocalState, Operator, OperatorOutput}; @@ -48,10 +48,7 @@ impl Operator for Filter { if input.num_rows() == 0 { return Ok(OperatorOutput::NeedMoreInput(None)); } - let partition = local - .downcast_ref::() - .map(|s| s.partition) - .unwrap_or(0); + let partition = local_mut::(local, "Filter")?.partition; let mask = self.predicate.evaluate_in(input.data(), partition)?; let out = input.filter(&mask)?; Ok(OperatorOutput::NeedMoreInput( diff --git a/crates/krishiv-exec/src/operators/project.rs b/crates/krishiv-exec/src/operators/project.rs index d29758a3e..27ce49145 100644 --- a/crates/krishiv-exec/src/operators/project.rs +++ b/crates/krishiv-exec/src/operators/project.rs @@ -8,8 +8,8 @@ use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions}; use arrow_schema::{Field, Schema, SchemaRef}; use krishiv_types::batch::Batch; -use super::PartitionState; use super::expr::ScalarExpr; +use super::{PartitionState, local_mut}; use crate::context::TaskContext; use crate::error::{Result, internal}; use crate::operator::{LocalState, Operator, OperatorOutput}; @@ -149,10 +149,7 @@ impl Operator for Project { if rows == 0 { return Ok(OperatorOutput::NeedMoreInput(None)); } - let partition = local - .downcast_ref::() - .map(|s| s.partition) - .unwrap_or(0); + let partition = local_mut::(local, "Project")?.partition; let data = input.data(); let columns = self .items diff --git a/crates/krishiv-exec/src/operators/tests.rs b/crates/krishiv-exec/src/operators/tests.rs index 744e53df5..c90dd9fa6 100644 --- a/crates/krishiv-exec/src/operators/tests.rs +++ b/crates/krishiv-exec/src/operators/tests.rs @@ -126,8 +126,9 @@ fn filter_keeps_selected_rows_with_their_weights() { #[test] fn filter_edge_cases_produce_no_output_and_null_does_not_select() { + let (_root, ctx) = ctx(); let f = Filter::new(schema(), gt1()); - let mut local = unlocal(); + let mut local = f.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); // Empty batch. let empty = ints(&[]); assert!( @@ -182,7 +183,11 @@ fn filter_errors_propagate() { Ok(BooleanArray::from(vec![true])) })), ); - let e = bad_len.execute(&ints(&[1, 2]), &mut unlocal()).unwrap_err(); + let (_root, ctx) = ctx(); + let mut local = bad_len + .local_state(&ctx) + .unwrap_or_else(|e| unreachable!("{e}")); + let e = bad_len.execute(&ints(&[1, 2]), &mut local).unwrap_err(); assert_eq!(e.error_class(), ErrorClass::InternalError); let failing = Filter::new( schema(), @@ -192,7 +197,7 @@ fn filter_errors_propagate() { })) })), ); - assert!(failing.execute(&ints(&[1]), &mut unlocal()).is_err()); + assert!(failing.execute(&ints(&[1]), &mut local).is_err()); assert!( format!( "{:?}", @@ -254,9 +259,11 @@ fn project_reorders_renames_and_computes() { .map(|f| f.name().clone()) .collect(); assert_eq!(names, ["s", "x", "y", "s"]); + let (_root, ctx) = ctx(); + let mut local = p.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); let input = weighted(&[1, 2], &[3, -4]); let out = out_batch( - p.execute(&input, &mut unlocal()) + p.execute(&input, &mut local) .unwrap_or_else(|e| unreachable!("{e}")), ) .unwrap_or_else(|| unreachable!("rows")); @@ -287,14 +294,16 @@ fn project_reorders_renames_and_computes() { #[test] fn project_edge_cases() { + let (_root, ctx) = ctx(); let all_null = batch(&[None, None]); let p = Project::try_new( &schema(), vec![ProjectItem::column(0), ProjectItem::expr("y", plus_one())], ) .unwrap_or_else(|e| unreachable!("{e}")); + let mut local = p.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); let out = out_batch( - p.execute(&all_null, &mut unlocal()) + p.execute(&all_null, &mut local) .unwrap_or_else(|e| unreachable!("{e}")), ) .unwrap_or_else(|| unreachable!("rows")); @@ -303,15 +312,18 @@ fn project_edge_cases() { // Empty input: no output. assert!( out_batch( - p.execute(&ints(&[]), &mut unlocal()) + p.execute(&ints(&[]), &mut local) .unwrap_or_else(|e| unreachable!("{e}")) ) .is_none() ); // No columns: the row count survives. let none = Project::try_new(&schema(), vec![]).unwrap_or_else(|e| unreachable!("{e}")); + let mut none_local = none + .local_state(&ctx) + .unwrap_or_else(|e| unreachable!("{e}")); let out = out_batch( - none.execute(&ints(&[1, 2, 3]), &mut unlocal()) + none.execute(&ints(&[1, 2, 3]), &mut none_local) .unwrap_or_else(|e| unreachable!("{e}")), ); assert_eq!( @@ -322,6 +334,7 @@ fn project_edge_cases() { #[test] fn project_rejects_bad_plans_and_bad_evaluators() { + let (_root, ctx) = ctx(); let e = Project::try_new(&schema(), vec![ProjectItem::column(2)]).unwrap_err(); assert_eq!(e.error_class(), ErrorClass::InternalError); // An evaluator whose column is shorter than the batch, or of another type. @@ -331,7 +344,8 @@ fn project_rejects_bad_plans_and_bad_evaluators() { vec![ProjectItem::expr("y", Arc::new(PlusOne { ty, short }))], ) .unwrap_or_else(|e| unreachable!("{e}")); - let e = p.execute(&ints(&[1, 2]), &mut unlocal()).unwrap_err(); + let mut local = p.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); + let e = p.execute(&ints(&[1, 2]), &mut local).unwrap_err(); assert_eq!(e.error_class(), ErrorClass::InternalError); } // An index that is valid for the plan but not for the batch. @@ -345,7 +359,22 @@ fn project_rejects_bad_plans_and_bad_evaluators() { ) .unwrap_or_else(|e| unreachable!("{e}")), ); - assert!(p.execute(&narrow, &mut unlocal()).is_err()); + let mut local = p.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); + assert!(p.execute(&narrow, &mut local).is_err()); +} + +#[test] +fn filter_and_project_reject_foreign_local_state() { + // A stateful operator evaluated with another operator's state fails + // closed instead of silently reading partition 0's stream. + let e = Filter::new(schema(), gt1()) + .execute(&ints(&[2]), &mut unlocal()) + .unwrap_err(); + assert_eq!(e.error_class(), ErrorClass::InternalError); + let p = Project::try_new(&schema(), vec![ProjectItem::column(0)]) + .unwrap_or_else(|e| unreachable!("{e}")); + let e = p.execute(&ints(&[2]), &mut unlocal()).unwrap_err(); + assert_eq!(e.error_class(), ErrorClass::InternalError); } fn run_limit(offset: u64, limit: u64, sizes: &[usize]) -> (Vec, usize) { @@ -967,8 +996,9 @@ fn expand_rejects_bad_plans() { } /// A stateful function sees the driver's partition: `rand(0)` draws -/// partition 0's stream on one driver and partition 1's on another, and -/// the same partition draws the same stream on a fresh evaluation. +/// partition 0's stream on one driver and partition 1's on another, even +/// when partition 1 binds first, and the same partition draws the same +/// stream on a fresh evaluation. #[test] fn project_threads_the_drivers_partition_to_stateful_functions() { use arrow_array::NullArray; @@ -1029,8 +1059,7 @@ fn project_threads_the_drivers_partition_to_stateful_functions() { ) .unwrap_or_else(|e| unreachable!("{e}")) }; - let run = |seed: i32, partition: u32| { - let p = project(seed); + let run = |p: &Arc, partition: u32| { let (_root, ctx) = ctx(); let ctx = ctx.with_partition(partition); let mut local = p.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); @@ -1041,10 +1070,13 @@ fn project_threads_the_drivers_partition_to_stateful_functions() { .unwrap_or_else(|| unreachable!("a row")); out.data().column(0).as_primitive::().value(0) }; - // Seed 0, partition 0 and 1: Spark's first draws. - assert_eq!(run(0, 0), 0.7604953758285915); - assert_eq!(run(0, 1), 0.6363787615254752); + // Partition 1 binds first on the shared binding: partition 0 still + // draws its own stream (each slot draws `base + slot`), and both are + // Spark's first draws for their seeds. + let p = project(0); + assert_eq!(run(&p, 1), 0.6363787615254752); + assert_eq!(run(&p, 0), 0.7604953758285915); // A fresh evaluation draws the same stream again: the stream is a pure // function of (seed, partition, row). - assert_eq!(run(0, 0), 0.7604953758285915); + assert_eq!(run(&project(0), 0), 0.7604953758285915); } diff --git a/crates/krishiv-functions/src/lib.rs b/crates/krishiv-functions/src/lib.rs index bd7e99c97..1720f815b 100644 --- a/crates/krishiv-functions/src/lib.rs +++ b/crates/krishiv-functions/src/lib.rs @@ -9,6 +9,9 @@ //! [`scalar::ScalarFunction`] over the `krishiv-kernels` kernels, which //! the physical planner calls once per batch. //! +//! - `uniform`: the `uniform` result-type rule both layers share (not a +//! layer: the analyzer and the executor map their own types to its kinds). +//! //! Aggregates are executed by `krishiv-exec`; the canonical operators and //! `CAST` by the expression evaluator. @@ -17,3 +20,4 @@ pub mod registry; pub mod scalar; pub mod signature; +pub(crate) mod uniform; diff --git a/crates/krishiv-functions/src/registry/misc.rs b/crates/krishiv-functions/src/registry/misc.rs index f294fc8e9..d4bfba31a 100644 --- a/crates/krishiv-functions/src/registry/misc.rs +++ b/crates/krishiv-functions/src/registry/misc.rs @@ -27,6 +27,7 @@ use krishiv_common::error::{ }; use krishiv_types::DataType; +use super::super::uniform::{UniformKind, uniform_result}; use super::defs::{Spec, scalar}; use crate::signature::{ArgType, CallTypes, Constant, Inputs}; @@ -143,29 +144,25 @@ fn is_uniform_bound(t: Option<&DataType>) -> bool { }) } -/// The wider of two integral types (Spark's `UpCastRule` order). -fn wider_integral(left: &DataType, right: &DataType) -> DataType { - let rank = |t: &DataType| match t { - DataType::Byte => 0, - DataType::Short => 1, - DataType::Integer => 2, - _ => 3, - }; - if rank(left) >= rank(right) { - left.clone() - } else { - right.clone() +/// The [`UniformKind`] of a `uniform` bound (an untyped `NULL` reads as +/// `DOUBLE`, as Spark casts it). +fn uniform_kind(t: Option<&DataType>) -> UniformKind { + match t { + None | Some(DataType::Double) => UniformKind::Double, + Some(DataType::Byte) => UniformKind::Byte, + Some(DataType::Short) => UniformKind::Short, + Some(DataType::Integer) => UniformKind::Int, + Some(DataType::Long) => UniformKind::Long, + Some(DataType::Float) => UniformKind::Float, + Some(DataType::Decimal(d)) => UniformKind::Decimal { + precision: d.precision(), + // A decimal scale is `0..=38`: the widening never fires. + scale: i8::try_from(d.scale()).unwrap_or(i8::MAX), + }, + Some(_) => UniformKind::Other, } } -/// Whether `left` is wider than `right` (Spark's -/// `DecimalType.isWiderThan`). -fn is_wider_decimal(left: (u8, u8), right: (u8, u8)) -> bool { - let (p1, s1) = (i32::from(left.0), i32::from(left.1)); - let (p2, s2) = (i32::from(right.0), i32::from(right.1)); - p1 - s1 >= p2 - s2 && s1 >= s2 && p1 >= p2 -} - /// `uniform(min, max[, seed])`: Spark's `Uniform.dataType` (an untyped /// `NULL` reads as `DOUBLE`), nullable when an argument may be. A NULL /// bound is a NULL result; a NULL seed reads as 0. @@ -210,32 +207,24 @@ fn uniform(c: &CallTypes<'_>) -> Result { Some(_) => {} } } - let num = |t: Option<&DataType>| t.cloned().unwrap_or(DataType::Double); - let (min, max) = (num(c.data_type(0)), num(c.data_type(1))); - let data_type = match (&min, &max) { - (a, b) - if matches!( - a, - DataType::Byte | DataType::Short | DataType::Integer | DataType::Long - ) && matches!( - b, - DataType::Byte | DataType::Short | DataType::Integer | DataType::Long - ) => - { - wider_integral(a, b) - } - (DataType::Double, _) | (_, DataType::Double) => DataType::Double, - (DataType::Float, _) | (_, DataType::Float) => DataType::Float, - (DataType::Decimal(left), DataType::Decimal(right)) - if is_wider_decimal( - (left.precision(), left.scale()), - (right.precision(), right.scale()), - ) => - { - min + let data_type = match uniform_result(uniform_kind(c.data_type(0)), uniform_kind(c.data_type(1))) + { + UniformKind::Byte => DataType::Byte, + UniformKind::Short => DataType::Short, + UniformKind::Int => DataType::Integer, + UniformKind::Long => DataType::Long, + UniformKind::Float => DataType::Float, + UniformKind::Double | UniformKind::Other => DataType::Double, + UniformKind::Decimal { precision, scale } => { + use krishiv_types::DecimalType; + // The kind came from a valid decimal bound, so the + // reconstruction holds; `try_new` guards the impossible fall. + let scale = u8::try_from(scale).unwrap_or(u8::MAX); + DataType::Decimal(DecimalType::try_new( + u32::from(precision), + u32::from(scale), + )?) } - (_, d @ DataType::Decimal(_)) | (d @ DataType::Decimal(_), _) => d.clone(), - _ => DataType::Double, }; Ok(ArgType::new(data_type, c.any_nullable())) } diff --git a/crates/krishiv-functions/src/registry/tests.rs b/crates/krishiv-functions/src/registry/tests.rs index 18c766401..9b8da1a5c 100644 --- a/crates/krishiv-functions/src/registry/tests.rs +++ b/crates/krishiv-functions/src/registry/tests.rs @@ -983,6 +983,56 @@ fn uniform_and_randstr_signatures() { assert!(err.to_string().contains("\"INT\""), "{err}"); } +/// The analyzer's `uniform` result type and the executor's bind type agree +/// on every pair of bound kinds: both layers map their own type to the +/// shared `uniform_result` rule, and the test pins the two mappers against +/// each other so one side cannot drift without the other. +#[test] +fn uniform_result_agrees_with_the_executor_bind() { + use crate::scalar::{BindArg, BindOptions, bind_scalar}; + // One foldable constant per bound type, as the analyzer sees them. + let cases: Vec<(Option, Option)> = vec![ + (Some(DataType::Byte), Some(Constant::Integral(1))), + (Some(DataType::Short), Some(Constant::Integral(1))), + (Some(DataType::Integer), Some(Constant::Integral(10))), + (Some(DataType::Long), Some(Constant::Integral(20))), + (Some(DataType::Float), Some(Constant::Fractional(1))), + (Some(DataType::Double), Some(Constant::Fractional(1))), + (Some(dec(10, 3)), Some(Constant::Fractional(1))), + (Some(dec(11, 4)), Some(Constant::Fractional(1))), + (None, Some(Constant::Null)), + ]; + for (a, ca) in &cases { + for (b, cb) in &cases { + let args = [ + ArgType { + data_type: a.clone(), + nullable: a.is_none(), + }, + ArgType { + data_type: b.clone(), + nullable: b.is_none(), + }, + ]; + let reg = result("uniform", &args, &[ca.clone(), cb.clone()]).unwrap(); + let bound = bind_scalar( + "uniform", + &[ + BindArg::Column(a.clone().unwrap_or(DataType::Null).to_arrow()), + BindArg::Column(b.clone().unwrap_or(DataType::Null).to_arrow()), + ], + &BindOptions::new(), + ) + .unwrap(); + assert_eq!( + reg.data_type.clone().map(|t| t.to_arrow()), + Some(bound.return_type().clone()), + "uniform({a:?}, {b:?})" + ); + } + } +} + #[test] fn bitwise_signatures() { let (b, s, i, l) = ( diff --git a/crates/krishiv-functions/src/scalar/misc.rs b/crates/krishiv-functions/src/scalar/misc.rs index ef7d3ce42..b31bc0e51 100644 --- a/crates/krishiv-functions/src/scalar/misc.rs +++ b/crates/krishiv-functions/src/scalar/misc.rs @@ -16,6 +16,11 @@ //! context yet, so they return what Spark returns with no file being read //! (the empty string, `-1`). //! +//! No golden pins the `version`, `input_file_*` or `uuid` values (no such +//! value queries in any `.out` file: `string-functions.sql` only pins +//! `length(uuid())` and that distinct calls differ), so those stand on the +//! documented behavior, unverified against a golden run. +//! //! `monotonically_increasing_id` is `(partition << 33) + row_in_partition` //! and `spark_partition_id` is the partition index, both from the //! executing partition (see [`super::ScalarFunction::invoke_in`]); each @@ -101,12 +106,16 @@ fn input_file_block() -> Binding { /// `spark_partition_id()`: the executing partition's index. fn spark_partition_id() -> Binding { Binding::partitioned(DataType::Int32, |x: &Inputs<'_>, partition| { + // Driver counts never approach `2^31` (one slot per worker thread). + debug_assert!(i32::try_from(partition).is_ok()); Ok(Arc::new(Int32Array::from(vec![partition as i32; x.rows()])) as ArrayRef) }) } /// `monotonically_increasing_id()`: `(partition << 33) + row_in_partition`, -/// counted per partition in the partitioned body. +/// counted per partition in the partitioned body. The per-partition counts +/// are untracked by `krishiv-memory`, but bounded by the partition count +/// (bytes), like the other per-partition state. fn monotonically_increasing_id() -> Binding { let counts = Mutex::new(Vec::::new()); Binding::partitioned(DataType::Int64, move |x: &Inputs<'_>, partition| { diff --git a/crates/krishiv-functions/src/scalar/random.rs b/crates/krishiv-functions/src/scalar/random.rs index 0495120ae..1401c76f3 100644 --- a/crates/krishiv-functions/src/scalar/random.rs +++ b/crates/krishiv-functions/src/scalar/random.rs @@ -19,7 +19,8 @@ //! exactly (`rand(0)` is `0.7604953758285915`). `uuid` draws its 128 bits //! from the same generator and sets the RFC 4122 version and variant bits; //! Spark fills those bits from a Mersenne Twister instead, so only the -//! shape (and the per-seed determinism) matches, not Spark's values. +//! shape (and the per-seed determinism) matches, not Spark's values. No +//! golden pins seeded `uuid` values (see `misc`). use std::sync::{Arc, Mutex}; @@ -30,6 +31,7 @@ use arrow_schema::DataType; use krishiv_common::error::KrishivError; use krishiv_kernels::cast::{CastMode, CastOptions, TimeZoneRules, cast}; +use super::super::uniform::{UniformKind, uniform_result as shared_result}; use super::args::{Inputs, internal}; use super::{BindArg, BindOptions, Binding, FamilyResult, bad_types, listed}; @@ -194,17 +196,34 @@ impl Stream { } /// The per-partition streams of one bound function, by partition index. +/// The seed is fixed per binding (the folded seed argument, or the query +/// seed for an absent or NULL seed), so the first bind records it and every +/// slot `j` draws `base + j`: a driver that binds a higher partition first +/// no longer seeds the lower slots with the wrong seed. The table is +/// untracked by `krishiv-memory`, but bounded by the partition count +/// (bytes), like the other per-partition counters. #[derive(Debug, Default)] struct Streams { + base: Option, streams: Vec, } impl Streams { - /// The stream of `partition`, seeded once. + /// The stream of `partition`: slot `j` draws `seed + j`, seeded once. + /// (`seed` is the binding's base seed; the callers pass it unfolded.) fn stream(&mut self, partition: u32, seed: i64) -> &mut Stream { + let base = *self.base.get_or_insert(seed); let index = partition as usize; if self.streams.len() <= index { - self.streams.resize_with(index + 1, || Stream::new(seed)); + let have = self.streams.len(); + self.streams.reserve(index + 1 - have); + for slot in have..=index { + // `slot <= index`, and `index` is a `u32` widened to + // `usize`: `slot` fits a `u32` (the fallback never fires). + let p = u32::try_from(slot).unwrap_or(u32::MAX); + self.streams + .push(Stream::new(base.wrapping_add(i64::from(p)))); + } } &mut self.streams[index] } @@ -302,7 +321,7 @@ fn rand(opts: &BindOptions, args: &[BindArg]) -> Binding { queries }; let mut streams = lock(&state)?; - let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let s = streams.stream(partition, seed); let out: Vec = (0..x.rows()).map(|_| s.rng.next_double()).collect(); Ok(Arc::new(Float64Array::from(out)) as ArrayRef) }) @@ -320,7 +339,7 @@ fn randn(opts: &BindOptions, args: &[BindArg]) -> Binding { queries }; let mut streams = lock(&state)?; - let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let s = streams.stream(partition, seed); let out: Vec = (0..x.rows()) .map(|_| s.rng.next_gaussian(&mut s.gaussian)) .collect(); @@ -341,7 +360,7 @@ fn uuid(opts: &BindOptions, args: &[BindArg]) -> Binding { queries }; let mut streams = lock(&state)?; - let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let s = streams.stream(partition, seed); let out: Vec = (0..x.rows()) .map(|_| { let most = (s.rng.next_long().cast_unsigned() & 0xffff_ffff_ffff_0fff) @@ -362,64 +381,38 @@ fn uuid(opts: &BindOptions, args: &[BindArg]) -> Binding { }) } -/// Whether `t` is an integral Arrow type. -fn is_integral(t: &DataType) -> bool { - matches!( - t, - DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 - ) -} - -/// The wider of two integral Arrow types (Spark's `UpCastRule` order). -fn wider_integral(a: &DataType, b: &DataType) -> DataType { - let rank = |t: &DataType| match t { - DataType::Int8 => 0, - DataType::Int16 => 1, - DataType::Int32 => 2, - _ => 3, - }; - if rank(a) >= rank(b) { - a.clone() - } else { - b.clone() +/// The [`UniformKind`] of a `uniform` bound (an untyped `NULL` reads as +/// `DOUBLE`, as Spark casts it). +fn uniform_kind(t: &DataType) -> UniformKind { + match t { + DataType::Int8 => UniformKind::Byte, + DataType::Int16 => UniformKind::Short, + DataType::Int32 => UniformKind::Int, + DataType::Int64 => UniformKind::Long, + DataType::Float32 => UniformKind::Float, + DataType::Float64 | DataType::Null => UniformKind::Double, + DataType::Decimal128(p, s) => UniformKind::Decimal { + precision: *p, + scale: *s, + }, + _ => UniformKind::Other, } } /// The Arrow result type of `uniform(min, max)` for argument types -/// `min`/`max` (an untyped `NULL` reads as `DOUBLE`, as Spark casts it): -/// Spark's `Uniform.dataType`. +/// `min`/`max`: Spark's `Uniform.dataType` (see `uniform_kind`). fn uniform_result(min: &DataType, max: &DataType) -> DataType { - let num = |t: &DataType| { - if *t == DataType::Null { - DataType::Float64 - } else { - t.clone() - } - }; - let (min, max) = (num(min), num(max)); - match (&min, &max) { - (a, b) if is_integral(a) && is_integral(b) => wider_integral(a, b), - (DataType::Float64, _) | (_, DataType::Float64) => DataType::Float64, - (DataType::Float32, _) | (_, DataType::Float32) => DataType::Float32, - (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => { - if is_wider_decimal((*p1, *s1), (*p2, *s2)) { - min - } else { - max - } - } - (_, d @ DataType::Decimal128(..)) | (d @ DataType::Decimal128(..), _) => d.clone(), - _ => DataType::Float64, + match shared_result(uniform_kind(min), uniform_kind(max)) { + UniformKind::Byte => DataType::Int8, + UniformKind::Short => DataType::Int16, + UniformKind::Int => DataType::Int32, + UniformKind::Long => DataType::Int64, + UniformKind::Float => DataType::Float32, + UniformKind::Double | UniformKind::Other => DataType::Float64, + UniformKind::Decimal { precision, scale } => DataType::Decimal128(precision, scale), } } -/// Spark's `DecimalType.isWiderThan`: every part at least as wide. -fn is_wider_decimal(left: (u8, i8), right: (u8, i8)) -> bool { - let (p1, s1) = (i32::from(left.0), i32::from(left.1)); - let (p2, s2) = (i32::from(right.0), i32::from(right.1)); - p1 - s1 >= p2 - s2 && s1 >= s2 && p1 >= p2 -} - /// `uniform(min, max[, seed])`: Spark's `Uniform` is a `RuntimeReplaceable` /// over `cast(min + (max - min) * Rand(seed))`, computed here in one pass: /// the bounds as doubles, one draw per row, then Spark's `Cast` to the @@ -443,7 +436,7 @@ fn uniform(opts: &BindOptions, args: &[BindArg]) -> Binding { return Ok(arrow_array::new_null_array(&out_type, x.rows())); }; let mut streams = lock(&state)?; - let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let s = streams.stream(partition, seed); let draws: Vec = (0..x.rows()) .map(|_| lo + (hi - lo) * s.rng.next_double()) .collect(); @@ -481,7 +474,7 @@ fn randstr(opts: &BindOptions, args: &[BindArg]) -> Binding { .into()); } let mut streams = lock(&state)?; - let s = streams.stream(partition, seed.wrapping_add(i64::from(partition))); + let s = streams.stream(partition, seed); let out: Vec = (0..x.rows()) .map(|_| { (0..n) diff --git a/crates/krishiv-functions/src/scalar/tests/random.rs b/crates/krishiv-functions/src/scalar/tests/random.rs index 60b3efeb7..ba2590b7a 100644 --- a/crates/krishiv-functions/src/scalar/tests/random.rs +++ b/crates/krishiv-functions/src/scalar/tests/random.rs @@ -118,6 +118,21 @@ fn partitions_own_their_streams() { assert_ne!(f64s(&a), f64s(&b)); } +#[test] +fn a_late_partition_does_not_reuse_an_early_partitions_seed() { + // Regression: a driver that bound a higher partition first seeded the + // lower slots with its own seed. Each slot draws `base + slot`, so + // partition 0 still draws Spark's first draw of seed 0 here. + let s = seed_i32(0); + let bound = vec![BindArg::Literal(Arc::clone(&s))]; + let f = super::bind_scalar("rand", &bound, &BindOptions::new()).unwrap(); + let a = lit(&s); + let p1 = f64s(&f.invoke_in(std::slice::from_ref(&a), 1, 1).unwrap()); + let p0 = f64s(&f.invoke_in(std::slice::from_ref(&a), 1, 0).unwrap()); + assert_eq!(p1, vec![Some(0.6363787615254752)]); + assert_eq!(p0, vec![Some(0.7604953758285915)]); +} + #[test] fn uniform_matches_the_golden_values() { // `SELECT uniform(0, 10, 0)` is 7, `(10, 20, 0)` is 17. diff --git a/crates/krishiv-functions/src/uniform.rs b/crates/krishiv-functions/src/uniform.rs new file mode 100644 index 000000000..8e1dbc0ca --- /dev/null +++ b/crates/krishiv-functions/src/uniform.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The `uniform(min, max)` result-type rule, shared by the analyzer rule +//! (`registry::misc`, over [`krishiv_types::DataType`]) and the executor +//! binding (`scalar::random`, over Arrow types): Spark's `Uniform.dataType`. +//! +//! Both layers map their own type to a [`UniformKind`]; [`uniform_result`] +//! then picks the result kind, so the rank order, the `double`/`float` +//! priority and the decimal rule live once. The layers only map kinds back +//! to their own type. `uniform_result_agrees_with_the_executor_bind` (in +//! `registry::tests`) pins every pair of kinds end to end across both +//! layers, so a mapper cannot drift from the rule. + +/// The numeric kinds `uniform`'s result rule distinguishes: the four +/// integral widths in Spark's `UpCastRule` order, the two floating widths, a +/// decimal width, and everything else (strings, booleans, timestamps, ...), +/// which reads as `DOUBLE` — as does an untyped `NULL` bound, which Spark +/// casts to `DOUBLE`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum UniformKind { + Byte, + Short, + Int, + Long, + Float, + Double, + Decimal { precision: u8, scale: i8 }, + Other, +} + +/// Whether `kind` is one of the four integral widths. +fn is_integral(kind: UniformKind) -> bool { + matches!( + kind, + UniformKind::Byte | UniformKind::Short | UniformKind::Int | UniformKind::Long + ) +} + +/// The result kind of `uniform(min, max)` for bound kinds `min`/`max`. +pub(crate) fn uniform_result(min: UniformKind, max: UniformKind) -> UniformKind { + let rank = |kind: UniformKind| match kind { + UniformKind::Byte => 0, + UniformKind::Short => 1, + UniformKind::Int => 2, + _ => 3, + }; + match (min, max) { + (a, b) if is_integral(a) && is_integral(b) => { + if rank(a) >= rank(b) { + a + } else { + b + } + } + (UniformKind::Double, _) | (_, UniformKind::Double) => UniformKind::Double, + (UniformKind::Float, _) | (_, UniformKind::Float) => UniformKind::Float, + ( + UniformKind::Decimal { + precision: p1, + scale: s1, + }, + UniformKind::Decimal { + precision: p2, + scale: s2, + }, + ) => { + if is_wider_decimal((p1, s1), (p2, s2)) { + min + } else { + max + } + } + (d @ UniformKind::Decimal { .. }, _) | (_, d @ UniformKind::Decimal { .. }) => d, + _ => UniformKind::Double, + } +} + +/// Whether the decimal `(p1, s1)` is wider than `(p2, s2)` (Spark's +/// `DecimalType.isWiderThan`: every part at least as wide). +pub(crate) fn is_wider_decimal(left: (u8, i8), right: (u8, i8)) -> bool { + let (p1, s1) = (i32::from(left.0), i32::from(left.1)); + let (p2, s2) = (i32::from(right.0), i32::from(right.1)); + p1 - s1 >= p2 - s2 && s1 >= s2 && p1 >= p2 +} diff --git a/crates/krishiv-session/tests/misc_functions.rs b/crates/krishiv-session/tests/misc_functions.rs index 7dc367144..2cb4fc3d9 100644 --- a/crates/krishiv-session/tests/misc_functions.rs +++ b/crates/krishiv-session/tests/misc_functions.rs @@ -2,14 +2,15 @@ //! The misc-function goldens end to end: every `misc-functions.sql`, //! `random.sql` and `current_database_catalog.sql` statement that uses only //! the implemented functions matches Spark 4.2.0 exactly (schema, rows, -//! error class and parameters), and the nondeterministic functions are +//! error class and parameters), except the multi-row `VALUES` statements: +//! the engine runs those tiny scans in one partition while Spark CI splits +//! them over two, so the row multisets differ. The single-partition draws +//! are asserted here; their decomposition into per-partition streams is +//! pinned at the binding level (`scalar/tests/random.rs`: partition 0's +//! first `uniform(10, 20, 0)` draw is 17, partition 1's next two are 16 and +//! 15 — the golden multiset — and `randstr`'s partition-1 rows match the +//! golden's rows 2-3 exactly). The nondeterministic functions are otherwise //! deterministic per partition. -//! -//! Two multi-row statements depend on Spark CI's partitioning (two slices -//! for three `VALUES` rows, then sorted output): the engine runs tiny scans -//! in one partition, so those exact row orders are asserted sorted where the -//! value multiset agrees, and at the operator level otherwise (see the -//! `#[test]` docs). #![cfg(test)] use arrow_array::cast::AsArray; @@ -208,6 +209,10 @@ fn typeof_needs_no_execution() { rows(&s, "select typeof(1 / 0)").0, vec!["typeof((CAST(1 AS DOUBLE) / CAST(0 AS DOUBLE)))".to_owned()] ); + assert_eq!( + rows(&s, "select typeof(1 / 0)").1, + vec![vec!["double".to_owned()]] + ); } #[test] @@ -238,8 +243,11 @@ fn current_database_and_catalog() { ); } +/// Seeded draws match Spark exactly; the multi-row `VALUES` cases assert +/// the engine's single-partition stream (Spark CI splits those rows over +/// two partitions — see the module docs). #[test] -fn rand_and_randn_match_spark() { +fn rand_and_randn_match_spark_single_partition() { let s = session(); // `random.sql`, exactly. assert_eq!( @@ -293,10 +301,16 @@ fn rand_and_randn_match_spark() { assert_eq!(e.error_class(), ErrorClass::SeedExpressionIsUnfoldable); let e = error(&s, "SELECT rand(id) FROM (SELECT 1 AS id)"); assert_eq!(e.error_class(), ErrorClass::SeedExpressionIsUnfoldable); + // Too many seeds is an arity error. + let e = error(&s, "SELECT rand(1, 2)"); + assert_eq!(e.error_class(), ErrorClass::WrongNumArgsWithoutSuggestion); } +/// Seeded values match Spark exactly; the multi-row `VALUES` case asserts +/// the engine's single-partition draws (Spark CI splits those rows over +/// two partitions — see the module docs). #[test] -fn uniform_matches_spark() { +fn uniform_matches_spark_single_partition() { let s = session(); let value = |sql: &str| rows(&s, sql).1.into_iter().flatten().collect::>(); assert_eq!(value("SELECT uniform(0, 1, 0) AS result"), ["0"]); @@ -324,6 +338,7 @@ fn uniform_matches_spark() { ["17.604953758285916"] ); assert_eq!(value("SELECT uniform(-10L, 10L, 0) AS result"), ["5"]); + assert_eq!(value("SELECT uniform(-20L, -10L, 0) AS result"), ["-12"]); assert_eq!(value("SELECT uniform(-20L, -10L, -10) AS result"), ["-17"]); // NULL bounds are NULL results; a NULL seed reads as 0. assert_eq!(value("SELECT uniform(NULL, 1, 0) AS result"), ["~"]); @@ -347,6 +362,9 @@ fn uniform_matches_spark() { assert!(e.to_string().contains("`seed`"), "{e}"); let e = error(&s, "SELECT uniform(10) AS result"); assert_eq!(e.error_class(), ErrorClass::WrongNumArgsWithoutSuggestion); + let e = error(&s, "SELECT uniform(10, 20, 30, 40) AS result"); + assert_eq!(e.error_class(), ErrorClass::WrongNumArgsWithoutSuggestion); + assert!(e.to_string().contains("[2, 3]"), "{e}"); let e = error(&s, "SELECT uniform(10.0F, 20.0F, 0.0F) AS result"); assert_eq!( e.error_class(), @@ -359,6 +377,13 @@ fn uniform_matches_spark() { ErrorClass::DatatypeMismatchUnexpectedInputType ); assert!(e.to_string().contains("\"NUMERIC\""), "{e}"); + let e = error(&s, "SELECT uniform(0, 'def', 0) AS result"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(e.to_string().contains("\"NUMERIC\""), "{e}"); + assert!(e.to_string().contains("second"), "{e}"); // Several partitions draw several streams: the engine runs the tiny // scan in one partition, so the rows are one stream's first draws // (Spark CI splits the rows over two partitions and sorts). @@ -372,8 +397,11 @@ fn uniform_matches_spark() { assert_eq!(got, ["10", "15", "17"]); } +/// Seeded strings match Spark exactly; the multi-row `VALUES` case asserts +/// the engine's single-partition draws (Spark CI splits those rows over +/// two partitions — see the module docs). #[test] -fn randstr_matches_spark() { +fn randstr_matches_spark_single_partition() { let s = session(); let value = |sql: &str| rows(&s, sql).1.into_iter().flatten().collect::>(); assert_eq!(value("SELECT randstr(1, 0) AS result"), ["c"]); @@ -407,6 +435,14 @@ fn randstr_matches_spark() { assert_eq!(e.error_class(), ErrorClass::WrongNumArgsWithoutSuggestion); let e = error(&s, "SELECT randstr(-1, 0) AS result"); assert_eq!(e.error_class(), ErrorClass::InvalidParameterValueLength); + // A string seed is neither an `int` nor a `bigint`. + let e = error(&s, "SELECT randstr(10, 'a') AS result"); + assert_eq!( + e.error_class(), + ErrorClass::DatatypeMismatchUnexpectedInputType + ); + assert!(e.to_string().contains("second"), "{e}"); + assert!(e.to_string().contains("(\"INT\" or \"BIGINT\")"), "{e}"); // The engine's tiny scan is one partition: consecutive draws of one // stream (Spark CI splits the rows over two partitions and sorts). let v = values_session(); From 415a21a78315deb291582ce23076deab0eff1f3a Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 02:11:55 +0100 Subject: [PATCH 3/7] functions: typeof is never NULL; randstr refuses an output past the i32 offset limit typeof is a literal of the argument's type name, so its result is non-nullable whatever the argument. randstr checks length times rows against the string array's i32 offsets before drawing or allocating and returns an error instead of panicking inside arrow. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-functions/src/registry/misc.rs | 6 +++-- .../krishiv-functions/src/registry/tests.rs | 6 +++++ crates/krishiv-functions/src/scalar/random.rs | 15 +++++++++++- .../src/scalar/tests/random.rs | 23 +++++++++++++++++++ 4 files changed, 47 insertions(+), 3 deletions(-) diff --git a/crates/krishiv-functions/src/registry/misc.rs b/crates/krishiv-functions/src/registry/misc.rs index d4bfba31a..3c2675780 100644 --- a/crates/krishiv-functions/src/registry/misc.rs +++ b/crates/krishiv-functions/src/registry/misc.rs @@ -50,8 +50,10 @@ fn int(c: &CallTypes<'_>) -> Result { /// `typeof(x)`: a non-nullable string. The name itself is resolved by the /// analyzer (a literal rewrite); the rule only types the call. -fn typeof_result(c: &CallTypes<'_>) -> Result { - string(c) +fn typeof_result(_: &CallTypes<'_>) -> Result { + // Spark's `typeof` is a literal of the argument's type name: never NULL, + // whatever the argument's nullability. + Ok(ArgType::new(DataType::String, false)) } /// Whether `t` is an `int` or a `bigint` (what every non-`uuid` seed takes). diff --git a/crates/krishiv-functions/src/registry/tests.rs b/crates/krishiv-functions/src/registry/tests.rs index d852325a0..3caba6663 100644 --- a/crates/krishiv-functions/src/registry/tests.rs +++ b/crates/krishiv-functions/src/registry/tests.rs @@ -1118,6 +1118,12 @@ fn misc_signatures() { assert_eq!(result("monotonically_increasing_id", &[], &[]).unwrap(), l); assert_eq!(result("spark_partition_id", &[], &[]).unwrap(), i); assert_eq!(result("typeof", std::slice::from_ref(&s), &[]).unwrap(), s); + // Never NULL, even for a nullable or untyped-NULL argument. + assert_eq!(result("typeof", &[null_ty()], &[]).unwrap(), s); + assert_eq!( + result("typeof", &[ArgType::new(DataType::String, true)], &[]).unwrap(), + s + ); assert_eq!(result("uuid", &[], &[]).unwrap(), s); assert_eq!(result("rand", &[], &[]).unwrap(), ty(DataType::Double)); assert_eq!( diff --git a/crates/krishiv-functions/src/scalar/random.rs b/crates/krishiv-functions/src/scalar/random.rs index 1401c76f3..b2044d2a7 100644 --- a/crates/krishiv-functions/src/scalar/random.rs +++ b/crates/krishiv-functions/src/scalar/random.rs @@ -28,7 +28,7 @@ use arrow_array::cast::AsArray; use arrow_array::types::{Float64Type, Int32Type, Int64Type}; use arrow_array::{ArrayRef, Float64Array, StringArray}; use arrow_schema::DataType; -use krishiv_common::error::KrishivError; +use krishiv_common::error::{KrishivError, UnableToAcquireMemory}; use krishiv_kernels::cast::{CastMode, CastOptions, TimeZoneRules, cast}; use super::super::uniform::{UniformKind, uniform_result as shared_result}; @@ -473,6 +473,19 @@ fn randstr(opts: &BindOptions, args: &[BindArg]) -> Binding { } .into()); } + // The output must fit a `StringArray`'s `i32` offsets: refuse before + // drawing or allocating, as the other string builders do. + let total = usize::try_from(n) + .ok() + .and_then(|n| n.checked_mul(x.rows())) + .filter(|t| i32::try_from(*t).is_ok()); + let Some(_) = total else { + return Err(UnableToAcquireMemory { + requested_bytes: n.to_string(), + received_bytes: "0".to_owned(), + } + .into()); + }; let mut streams = lock(&state)?; let s = streams.stream(partition, seed); let out: Vec = (0..x.rows()) diff --git a/crates/krishiv-functions/src/scalar/tests/random.rs b/crates/krishiv-functions/src/scalar/tests/random.rs index ba2590b7a..bb19d79b7 100644 --- a/crates/krishiv-functions/src/scalar/tests/random.rs +++ b/crates/krishiv-functions/src/scalar/tests/random.rs @@ -263,6 +263,29 @@ fn randstr_null_length_is_empty_and_negative_fails() { assert!(err.to_string().contains("-1"), "{err}"); } +#[test] +fn randstr_refuses_an_output_past_the_string_offset_limit() { + use krishiv_common::error::ErrorClass; + // 1e9 characters over 3 rows exceed `i32` offsets: an error, not a panic. + let (big, sd) = (seed_i32(1_000_000_000), seed_i32(0)); + let bound = vec![ + BindArg::Literal(Arc::clone(&big)), + BindArg::Literal(Arc::clone(&sd)), + ]; + let f = super::bind_scalar("randstr", &bound, &BindOptions::new()).unwrap(); + let err = f.invoke_in(&[lit(&big), lit(&sd)], 3, 0).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::UnableToAcquireMemory); + // The binding is still usable afterwards (nothing was poisoned). + let small = seed_i32(3); + let bound = vec![ + BindArg::Literal(Arc::clone(&small)), + BindArg::Literal(Arc::clone(&sd)), + ]; + let f = super::bind_scalar("randstr", &bound, &BindOptions::new()).unwrap(); + let out = f.invoke_in(&[lit(&small), lit(&sd)], 2, 0).unwrap(); + assert_eq!(out.as_string::().value(0), "ceV"); +} + #[test] fn uuids_have_shape_and_partitioned_streams() { let sd = seed_i32(42); From 1f46a8f795eefd0a9b247d5350f59513d39d2a85 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 02:35:46 +0100 Subject: [PATCH 4/7] testkit: typeof registry-matrix entries no longer differ from Spark typeof is now non-nullable, so the recorded nullability mismatches for its NULL arguments are gone. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../analysis/registry-matrix/typeof.json | 24 ------------------- 1 file changed, 24 deletions(-) diff --git a/tests/fixtures/analysis/registry-matrix/typeof.json b/tests/fixtures/analysis/registry-matrix/typeof.json index 84f7ab53e..71c123e10 100644 --- a/tests/fixtures/analysis/registry-matrix/typeof.json +++ b/tests/fixtures/analysis/registry-matrix/typeof.json @@ -8,14 +8,6 @@ "id": "typeof#737efe3cca2b/ansi_off", "reason": "SELECT typeof(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map" }, - { - "id": "typeof#7b81ac5ac6dd/ansi_on", - "reason": "SELECT typeof(CAST(NULL AS array)) | string where Spark has string not null" - }, - { - "id": "typeof#7b81ac5ac6dd/ansi_off", - "reason": "SELECT typeof(CAST(NULL AS array)) | string where Spark has string not null" - }, { "id": "typeof#85119c3e6111/ansi_on", "reason": "SELECT typeof(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals" @@ -64,22 +56,6 @@ "id": "typeof#bea7c4b2a1ca/ansi_off", "reason": "SELECT typeof(CAST(array(NULL) AS array)) | KRISHIV_UNSUPPORTED_FEATURE: function array" }, - { - "id": "typeof#e22394e72946/ansi_on", - "reason": "SELECT typeof(CAST(NULL AS int)) | string where Spark has string not null" - }, - { - "id": "typeof#e22394e72946/ansi_off", - "reason": "SELECT typeof(CAST(NULL AS int)) | string where Spark has string not null" - }, - { - "id": "typeof#e571d1babdbc/ansi_on", - "reason": "SELECT typeof(NULL) | string where Spark has string not null" - }, - { - "id": "typeof#e571d1babdbc/ansi_off", - "reason": "SELECT typeof(NULL) | string where Spark has string not null" - }, { "id": "typeof#ef1847dcc457/ansi_on", "reason": "SELECT typeof(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals" From 3d5c1a96e9661bf5a92d611d473ce2d0880482ba Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 05:28:10 +0100 Subject: [PATCH 5/7] test(functions): pin the random streams to a Spark 4.2.0 oracle rand_advances_one_draw_per_row and partitions_own_their_streams assert the exact first draws of partitions 0 and 1 for seed 0, as returned by SELECT spark_partition_id(), rand(0) FROM range(0, 6, 1, 2), and the second batch of a partition continues its stream. The stale comments about uniform and monotonically_increasing_id are gone. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../src/scalar/tests/random.rs | 49 ++++++++++++++----- 1 file changed, 36 insertions(+), 13 deletions(-) diff --git a/crates/krishiv-functions/src/scalar/tests/random.rs b/crates/krishiv-functions/src/scalar/tests/random.rs index bb19d79b7..535914f21 100644 --- a/crates/krishiv-functions/src/scalar/tests/random.rs +++ b/crates/krishiv-functions/src/scalar/tests/random.rs @@ -50,17 +50,18 @@ fn rand_zero_matches_spark() { #[test] fn rand_advances_one_draw_per_row() { - // `SELECT uniform(10, 20, 0) FROM VALUES (0), (1), (2)` is 15, 16, 17: - // the rows draw consecutively from one stream. + // Oracle: `SELECT rand(0) FROM range(0, 3, 1, 1)` on a Spark 4.2.0 + // server: the rows draw consecutively from one stream. let s = seed_i32(0); let got = seeded("rand", &s, 3, 0); - let rows = f64s(&got); - assert_eq!(rows.len(), 3); - assert!( - (rows[0].unwrap() - 0.7604953758285915).abs() < 1e-15, - "{rows:?}" + assert_eq!( + f64s(&got), + vec![ + Some(0.7604953758285915), + Some(0.5234194256885571), + Some(0.0953472826424725) + ] ); - assert!(rows.windows(2).all(|w| w[0] != w[1]), "{rows:?}"); } #[test] @@ -99,13 +100,31 @@ fn long_seeds_match_int_seeds() { fn partitions_own_their_streams() { // The same query in another partition draws another stream; the // partition seeds the generator (`seed + partitionIndex`, Spark's - // `initializeInternal`). + // `initializeInternal`). Oracle: `SELECT spark_partition_id(), rand(0) + // FROM range(0, 6, 1, 2)` on a Spark 4.2.0 server gives partition 0 + // 0.7604953758285915, 0.5234194256885571, 0.0953472826424725 and + // partition 1 0.6363787615254752, 0.5993846534021868, 0.134842710012538. let s = seed_i32(0); let p0 = f64s(&seeded("rand", &s, 3, 0)); let p1 = f64s(&seeded("rand", &s, 3, 1)); - assert_ne!(p0, p1); - // `monotonically_increasing_id` and `rand` of two partitions differ, - // and one partition is stable across batches. + assert_eq!( + p1, + vec![ + Some(0.6363787615254752), + Some(0.5993846534021868), + Some(0.134842710012538) + ] + ); + assert_eq!( + p0, + vec![ + Some(0.7604953758285915), + Some(0.5234194256885571), + Some(0.0953472826424725) + ] + ); + // One partition continues its stream across batches: the second batch + // starts where the first stopped (the third draw of the stream). let f = super::bind_scalar( "rand", &[BindArg::Literal(Arc::clone(&s))], @@ -115,7 +134,11 @@ fn partitions_own_their_streams() { let lit = Arg::Scalar(Arc::clone(&s)); let a = f.invoke_in(std::slice::from_ref(&lit), 2, 0).unwrap(); let b = f.invoke_in(std::slice::from_ref(&lit), 2, 0).unwrap(); - assert_ne!(f64s(&a), f64s(&b)); + assert_eq!( + f64s(&a), + vec![Some(0.7604953758285915), Some(0.5234194256885571)] + ); + assert_eq!(f64s(&b)[0], Some(0.0953472826424725)); } #[test] From a8a825b5e31d6f2c08ffe70a9b48ca065f785de2 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:00:24 +0100 Subject: [PATCH 6/7] fix(exec): partition-aware evaluation in scans, joins and replicate; unique partition ids across fan-in members Scan pushed filters, hash, nested-loop and grace join residuals and the replicate count evaluate in the driver's partition (evaluate_in), so rand, uuid, monotonically_increasing_id and spark_partition_id read that partition's state instead of partition 0. The drivers of every member of a pipeline group share one partition id space, numbered member after member as Spark's UnionRDD numbers a union's partitions; an id past u32 is an error instead of wrapping. Pinned against a Spark 4.2.0 run: spark_partition_id over a three-way UNION ALL is 0, 1, 2 and monotonically_increasing_id over a two-way one 0, 8589934592. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/context.rs | 15 +++---- crates/krishiv-exec/src/driver.rs | 5 ++- crates/krishiv-exec/src/join/grace.rs | 36 +++++++++++----- crates/krishiv-exec/src/join/nested.rs | 10 ++++- crates/krishiv-exec/src/join/probe.rs | 15 +++++-- .../krishiv-exec/src/operators/replicate.rs | 5 ++- crates/krishiv-exec/src/operators/scan.rs | 16 +++++-- crates/krishiv-exec/src/operators/tests.rs | 3 +- crates/krishiv-exec/src/runner.rs | 42 +++++++++++++++---- .../krishiv-session/tests/misc_functions.rs | 17 ++++++++ 10 files changed, 126 insertions(+), 38 deletions(-) diff --git a/crates/krishiv-exec/src/context.rs b/crates/krishiv-exec/src/context.rs index 3c511db03..220c9c84e 100644 --- a/crates/krishiv-exec/src/context.rs +++ b/crates/krishiv-exec/src/context.rs @@ -398,9 +398,10 @@ pub struct TaskContext { name: &'static str, memory: OnceLock, dependency: Option, - /// The driver's index in its pipeline run (Spark's partition index for - /// the partition the driver computes): stateful functions own one value - /// stream per partition. One driver always reads partition 0. + /// The driver's scheduler-assigned globally unique partition id within + /// its pipeline group (Spark's partition index for the partition the + /// driver computes): stateful functions own one value stream per + /// partition. The first driver of a group reads partition 0. partition: u32, } @@ -434,16 +435,16 @@ impl TaskContext { } } - /// The same context for the driver `partition` (its index in the - /// pipeline run). + /// The same context for the driver `partition` (its scheduler-assigned + /// globally unique id within its pipeline group). #[must_use] pub(crate) fn with_partition(mut self, partition: u32) -> Self { self.partition = partition; self } - /// The driver's index in its pipeline run (set when the driver is - /// spawned; 0 for a single driver). + /// The driver's scheduler-assigned partition id (set when the driver is + /// spawned; the first driver of a group reads 0). pub fn partition(&self) -> u32 { self.partition } diff --git a/crates/krishiv-exec/src/driver.rs b/crates/krishiv-exec/src/driver.rs index 3bb1f8e05..7f7c7d600 100644 --- a/crates/krishiv-exec/src/driver.rs +++ b/crates/krishiv-exec/src/driver.rs @@ -119,8 +119,9 @@ fn since(executor: &dyn Executor, start: Instant) -> Duration { executor.now().saturating_duration_since(start) } -/// Spawns one driver of `run`: `partition` is the driver's index in the -/// run (see [`TaskContext::with_partition`]). +/// Spawns one driver of `run`: `partition` is the driver's +/// scheduler-assigned globally unique partition id within its pipeline +/// group (see [`TaskContext::with_partition`]). pub(crate) fn spawn_driver(run: Arc, partition: u32) -> TaskHandle { let pipeline = run.pipeline.index; let query = Arc::clone(&run.query); diff --git a/crates/krishiv-exec/src/join/grace.rs b/crates/krishiv-exec/src/join/grace.rs index 7606aeab1..d24d49514 100644 --- a/crates/krishiv-exec/src/join/grace.rs +++ b/crates/krishiv-exec/src/join/grace.rs @@ -516,8 +516,9 @@ async fn reverse( let hashing = hashing.clone(); let handle = Arc::clone(ctx.consumer().consumer()); let query = Arc::clone(ctx.query()); + let partition = ctx.partition(); let out = ctx - .run_blocking(move || reverse_blocking(&shared, &hashing, task, &handle, &query)) + .run_blocking(move || reverse_blocking(&shared, &hashing, task, &handle, &query, partition)) .await; ctx.consumer().consumer().release_credit(); out @@ -536,6 +537,7 @@ fn reverse_blocking( task: Task, handle: &Arc, query: &QueryContext, + partition: u32, ) -> Result { // Everything the join holds is reserved before it is allocated, and the // reversal declines while it still can (before the build stream starts): @@ -648,10 +650,10 @@ fn reverse_blocking( } } if pairs.len() >= PAIR_BATCH { - mark_pairs(shared, &probe, &mut pairs, &mut matched)?; + mark_pairs(shared, &probe, &mut pairs, &mut matched, partition)?; } } - mark_pairs(shared, &probe, &mut pairs, &mut matched) + mark_pairs(shared, &probe, &mut pairs, &mut matched, partition) }; for_each( build, @@ -684,6 +686,7 @@ fn mark_pairs( probe: &Batch, pairs: &mut Vec<(u32, &[u8])>, matched: &mut [bool], + partition: u32, ) -> Result<()> { let Some(filter) = &shared.filter else { pairs.clear(); @@ -701,7 +704,7 @@ fn mark_pairs( let mut columns: Vec = taken.data().columns().to_vec(); columns.extend(build_cols); let candidates = super::record(&shared.filter_schema, columns, m)?; - let mask = filter.evaluate(&candidates)?; + let mask = filter.evaluate_in(&candidates, partition)?; if mask.len() != m { return Err(internal( "a join filter returned the wrong row count".to_owned(), @@ -770,10 +773,11 @@ pub(crate) async fn step( .get() .ok_or_else(|| internal("hash join probed before its build".to_owned()))?; let limit = ctx.query().config().batch_size.max(1); + let partition = ctx.partition(); if let Some(a) = &mut st.active { if let Some(c) = &mut a.cursor { let tables = Tables::Single(&a.table); - let (out, done) = c.next(shared, &tables, limit)?; + let (out, done) = c.next(shared, &tables, limit, partition)?; if done { a.cursor = None; } @@ -1676,9 +1680,14 @@ mod tests { let mut filler = krishiv_memory::Reservation::new(Arc::clone(&other), id); filler.try_grow(LIMIT - root.reserved()).unwrap(); let handle = Arc::clone(ctx.consumer().consumer()); - let Ok(Reversed::Declined(task)) = - reverse_blocking(shared, &hashing, task, &handle, ctx.query()) - else { + let Ok(Reversed::Declined(task)) = reverse_blocking( + shared, + &hashing, + task, + &handle, + ctx.query(), + ctx.partition(), + ) else { panic!("the probe rows cannot fit"); }; assert_eq!(task.probe.files.len(), files, "nothing was consumed"); @@ -1686,9 +1695,14 @@ mod tests { // With room, the same task joins: a semi join with no build rows // emits nothing, and the probe files are gone with the task. drop(filler); - let Ok(Reversed::Joined(out)) = - reverse_blocking(shared, &hashing, task, &handle, ctx.query()) - else { + let Ok(Reversed::Joined(out)) = reverse_blocking( + shared, + &hashing, + task, + &handle, + ctx.query(), + ctx.partition(), + ) else { panic!("the probe rows fit now"); }; assert!(out.is_none()); diff --git a/crates/krishiv-exec/src/join/nested.rs b/crates/krishiv-exec/src/join/nested.rs index 441d6a587..511d2dbcc 100644 --- a/crates/krishiv-exec/src/join/nested.rs +++ b/crates/krishiv-exec/src/join/nested.rs @@ -221,6 +221,8 @@ struct ProbeLocal { offset: usize, /// Per probe row of the chunk: whether it matched. matched: Vec, + /// The driver's partition: the residual filter evaluates in it. + partition: u32, } impl Operator for ProbeOp { @@ -250,6 +252,7 @@ impl Operator for ProbeOp { batch: 0, offset: 0, matched: Vec::new(), + partition: ctx.partition(), })) } @@ -280,8 +283,10 @@ impl Operator for ProbeOp { let n = (b.num_rows() - l.offset).min((TARGET_PAIRS / (l.end - l.start)).max(1)); let slice = b.slice(l.offset, n); l.offset += n; + let (partition, start, end) = (l.partition, l.start, l.end); + let matched = &mut l.matched; if let Some(pairs) = - self.pairs(input.data(), l.start, l.end, &slice, &mut l.matched)? + self.pairs(input.data(), start, end, &slice, matched, partition)? { out.push(pairs); } @@ -325,6 +330,7 @@ impl ProbeOp { end: usize, build: &RecordBatch, matched: &mut [bool], + partition: u32, ) -> Result> { let b = build.num_rows(); if b == 0 { @@ -356,7 +362,7 @@ impl ProbeOp { let mask = match &self.0.spec.filter { None => BooleanArray::from(vec![true; n]), Some(f) => { - let m = f.evaluate(&paired)?; + let m = f.evaluate_in(&paired, partition)?; // NULL is not a match. match m.nulls() { Some(v) => BooleanArray::new(m.values() & v.inner(), None), diff --git a/crates/krishiv-exec/src/join/probe.rs b/crates/krishiv-exec/src/join/probe.rs index 7678d1cbc..226d5816f 100644 --- a/crates/krishiv-exec/src/join/probe.rs +++ b/crates/krishiv-exec/src/join/probe.rs @@ -154,6 +154,7 @@ impl Cursor { shared: &Shared, tables: &Tables<'_>, limit: usize, + partition: u32, ) -> Result<(Option, bool)> { let n = self.batch.num_rows(); let quick = !shared.join_type.emits_build() && shared.filter.is_none(); @@ -190,7 +191,10 @@ impl Cursor { } } if !probe.is_empty() { - return Ok((self.pairs(shared, tables, &probe, &build)?, false)); + return Ok(( + self.pairs(shared, tables, &probe, &build, partition)?, + false, + )); } Ok((self.tail(shared)?, true)) } @@ -203,6 +207,7 @@ impl Cursor { tables: &Tables<'_>, probe: &[u32], build: &[(usize, u32)], + partition: u32, ) -> Result> { let arena = |p: usize| tables.get(p).map(JoinTable::arena); let build_cols = shared.build_columns(self.rows, build, arena)?; @@ -213,7 +218,7 @@ impl Cursor { let mask = match &shared.filter { Some(f) => { let candidates = record(&shared.filter_schema, columns.clone(), m)?; - let mask = f.evaluate(&candidates)?; + let mask = f.evaluate_in(&candidates, partition)?; if mask.len() != m { return Err(internal( "a join filter returned the wrong row count".to_owned(), @@ -368,6 +373,8 @@ struct ProbeLocal { signal: Arc, flush: Flush, batch_size: usize, + /// The driver's partition: the residual filter evaluates in it. + partition: u32, } fn local_mut(local: &mut LocalState) -> Result<&mut ProbeLocal> { @@ -428,12 +435,13 @@ impl ProbeOp { l.cursor = Some(cursor); } let limit = l.batch_size; + let partition = l.partition; let cursor = l .cursor .as_mut() .ok_or_else(|| internal("no probe cursor".to_owned()))?; loop { - let (out, done) = cursor.next(s, &tables, limit)?; + let (out, done) = cursor.next(s, &tables, limit, partition)?; if done { l.cursor = None; return Ok(OperatorOutput::NeedMoreInput(out)); @@ -473,6 +481,7 @@ impl Operator for ProbeOp { signal: Arc::clone(memory.signal()), flush: Flush::default(), batch_size: ctx.query().config().batch_size.max(1), + partition: ctx.partition(), })) } diff --git a/crates/krishiv-exec/src/operators/replicate.rs b/crates/krishiv-exec/src/operators/replicate.rs index 3fd4395b2..e5aac7f24 100644 --- a/crates/krishiv-exec/src/operators/replicate.rs +++ b/crates/krishiv-exec/src/operators/replicate.rs @@ -38,6 +38,8 @@ impl ReplicateRows { /// Where the operator is in the current input batch. struct Cursor { batch_size: usize, + /// The driver's partition: `count` evaluates stateful functions in it. + partition: u32, /// The input's counts, while some of its copies are still to emit. counts: Option, /// The row being copied. @@ -58,6 +60,7 @@ impl Operator for ReplicateRows { fn local_state(&self, ctx: &TaskContext) -> Result { Ok(Box::new(Cursor { batch_size: ctx.query().config().batch_size.max(1), + partition: ctx.partition(), counts: None, row: 0, emitted: 0, @@ -72,7 +75,7 @@ impl Operator for ReplicateRows { if input.num_rows() == 0 { return Ok(OperatorOutput::NeedMoreInput(None)); } - let v = self.count.evaluate(input.data())?; + let v = self.count.evaluate_in(input.data(), cur.partition)?; let c = v .as_primitive_opt::() .ok_or_else(|| internal("a replication count is not BIGINT".to_owned()))? diff --git a/crates/krishiv-exec/src/operators/scan.rs b/crates/krishiv-exec/src/operators/scan.rs index 770f9b637..3663e332c 100644 --- a/crates/krishiv-exec/src/operators/scan.rs +++ b/crates/krishiv-exec/src/operators/scan.rs @@ -195,6 +195,9 @@ pub(super) struct ScanLocal { pub(super) claimed: u64, /// The runtime filters' row stage. pub(super) rows: super::RowStage, + /// The driver's partition: pushed filters evaluate stateful functions + /// in it (see [`Predicate::evaluate_in`]). + pub(super) partition: u32, } impl MemoryScan { @@ -278,7 +281,12 @@ impl MemoryScan { /// Reads the next piece of the claimed (non-empty) morsel, filtered and /// projected. - pub(super) fn next_piece(&self, local: &mut ScanLocal, batch_rows: usize) -> Result { + pub(super) fn next_piece( + &self, + local: &mut ScanLocal, + batch_rows: usize, + partition: u32, + ) -> Result { let batch = self .batches .get(local.batches.start) @@ -301,7 +309,7 @@ impl MemoryScan { if piece.num_rows() == 0 { break; } - let mask = filter.evaluate(piece.data())?; + let mask = filter.evaluate_in(piece.data(), partition)?; piece = piece.filter(&mask)?; } // After the pushed filters, as Spark evaluates a runtime filter @@ -352,7 +360,8 @@ impl MemoryScan { // One piece per call, even an empty input batch or one a pushed // filter emptied: the driver must reach a batch boundary (cancel, // abort and spill checks) after every piece read. - let piece = self.next_piece(local, g.batch_rows)?; + let partition = local.partition; + let piece = self.next_piece(local, g.batch_rows, partition)?; g.returned .fetch_add(piece.num_rows() as u64, Ordering::Relaxed); Ok(Some(piece)) @@ -421,6 +430,7 @@ impl Source for MemoryScan { offset: 0, claimed: 0, rows: super::RowStage::new(&g.runtime, &self.table, self.runtime.settings()), + partition: _ctx.partition(), })) } diff --git a/crates/krishiv-exec/src/operators/tests.rs b/crates/krishiv-exec/src/operators/tests.rs index c90dd9fa6..38342e8ce 100644 --- a/crates/krishiv-exec/src/operators/tests.rs +++ b/crates/krishiv-exec/src/operators/tests.rs @@ -740,8 +740,9 @@ fn scan_reports_inconsistent_state_as_internal_errors() { offset: 0, claimed: 0, rows: super::RowStage::new(&[], &schema(), super::RowStageSettings::default()), + partition: 0, }; - let e = scan.next_piece(&mut local, 4).unwrap_err(); + let e = scan.next_piece(&mut local, 4, 0).unwrap_err(); assert_eq!(e.error_class(), ErrorClass::InternalError); // A cursor that hands out a morsel the global state does not have. let global: crate::GlobalState = Arc::new(ScanGlobal { diff --git a/crates/krishiv-exec/src/runner.rs b/crates/krishiv-exec/src/runner.rs index 1bd714502..c3c982779 100644 --- a/crates/krishiv-exec/src/runner.rs +++ b/crates/krishiv-exec/src/runner.rs @@ -401,6 +401,12 @@ fn start_group( let (sink_id, sink) = &first.sink; let sink_global = sink.global_state(&ctx((*sink_id, sink.name())))?; let mut started = Vec::with_capacity(members.len()); + // Scheduler-assigned globally unique partition ids: the drivers of + // every member share one id space, in pipeline order, so two drivers + // never compute the same partition (stateful functions own one value + // stream per partition). An index past `u32` is rejected rather than + // wrapped to `u32::MAX`, which would collide. + let mut next_partition: u64 = 0; for pipeline in members { let (source_id, source) = &pipeline.source; let mut source_ctx = ctx((*source_id, source.name())); @@ -421,11 +427,18 @@ fn start_group( .min(source.max_parallelism(&source_global)) .max(1) }; + let base = u32::try_from(next_partition) + .map_err(|_| internal("too many drivers for u32 partition ids".to_owned()))?; + // `drivers` fits `usize`; the end index must fit `u32` as well. + u32::try_from(next_partition + drivers as u64) + .map_err(|_| internal("too many drivers for u32 partition ids".to_owned()))?; + next_partition += drivers as u64; started.push(MemberStart { pipeline, operator_deps, source_global, drivers, + base_partition: base, }); } // One barrier per operator, over the drivers of every pipeline it is in. @@ -455,6 +468,8 @@ struct MemberStart { operator_deps: Vec>, source_global: GlobalState, drivers: usize, + /// The first scheduler-assigned partition id of its drivers. + base_partition: u32, } /// A pipeline group planned but not yet running: the state its members @@ -495,6 +510,21 @@ fn member_barriers( .ok_or_else(|| internal("an operator without a barrier".to_owned())) } +/// One driver per index of `member`: the scheduler-assigned partition ids +/// `base..base + count`. Rejects an id past `u32` instead of wrapping it. +fn drivers(run: &Arc, base: u32, count: usize) -> Result> { + let mut out = Vec::with_capacity(count); + for i in 0..count { + let offset = u32::try_from(i) + .map_err(|_| internal("too many drivers for u32 partition ids".to_owned()))?; + let partition = base + .checked_add(offset) + .ok_or_else(|| internal("too many drivers for u32 partition ids".to_owned()))?; + out.push(spawn_driver(Arc::clone(run), partition)); + } + Ok(out) +} + /// Starts every member of `group` together, sharing the sink's state. fn spawn_group( executor: &Arc, @@ -522,10 +552,7 @@ fn spawn_group( query: Arc::clone(query), executor: Arc::clone(executor), }); - handles.extend((0..m.drivers).map(|i| { - let partition = u32::try_from(i).unwrap_or(u32::MAX); - spawn_driver(Arc::clone(&run), partition) - })); + handles.extend(drivers(&run, m.base_partition, m.drivers)?); runs.push(run); } Ok((runs, handles)) @@ -560,10 +587,9 @@ async fn run_group_sequential( executor: Arc::clone(executor), }); let mut handles = Drivers(FuturesUnordered::new()); - handles.0.extend((0..m.drivers).map(|i| { - let partition = u32::try_from(i).unwrap_or(u32::MAX); - spawn_driver(Arc::clone(&run), partition) - })); + handles + .0 + .extend(drivers(&run, m.base_partition, m.drivers)?); await_drivers(query, &mut handles).await; record_extra_metrics(query, std::slice::from_ref(&run)); query.check_cancelled()?; diff --git a/crates/krishiv-session/tests/misc_functions.rs b/crates/krishiv-session/tests/misc_functions.rs index 2cb4fc3d9..470e98d7a 100644 --- a/crates/krishiv-session/tests/misc_functions.rs +++ b/crates/krishiv-session/tests/misc_functions.rs @@ -483,6 +483,23 @@ fn nondeterministic_queries_are_deterministic_per_partition() { assert_eq!(value("SELECT input_file_name()"), [""]); assert_eq!(value("SELECT input_file_block_start()"), ["-1"]); assert_eq!(value("SELECT input_file_block_length()"), ["-1"]); + // UNION ALL numbers its children's partitions one after another, as + // Spark's `UnionRDD` does (Spark 4.2.0, local[1]: 0, 1, 2 and + // 0, 8589934592). + assert_eq!( + value( + "SELECT spark_partition_id() AS p FROM (SELECT 1 AS x UNION ALL \ + SELECT 2 AS x UNION ALL SELECT 3 AS x) ORDER BY p" + ), + ["0", "1", "2"] + ); + assert_eq!( + value( + "SELECT monotonically_increasing_id() AS m FROM (SELECT 1 AS x \ + UNION ALL SELECT 2 AS x) ORDER BY m" + ), + ["0", "8589934592"] + ); // A stateful predicate sees the same partition stream. let v = values_session(); assert_eq!( From 217077bd7cef46070509c364d5df37b432d5d2a4 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:26:53 +0100 Subject: [PATCH 7/7] test: refresh the registry-matrix manifest and golden ratchet against main named_struct now resolves, so the random functions' struct-argument entries leave the manifest; order-by#5, #6 and #12 now pass. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../analysis/registry-matrix/rand.json | 8 ------- .../analysis/registry-matrix/randn.json | 8 ------- .../analysis/registry-matrix/random.json | 8 ------- .../analysis/registry-matrix/randstr.json | 16 ------------- .../analysis/registry-matrix/typeof.json | 8 ------- .../analysis/registry-matrix/uniform.json | 24 ------------------- .../known-failures-cases/order-by.json | 5 ---- 7 files changed, 77 deletions(-) diff --git a/tests/fixtures/analysis/registry-matrix/rand.json b/tests/fixtures/analysis/registry-matrix/rand.json index f70faf2ce..c92cc4430 100644 --- a/tests/fixtures/analysis/registry-matrix/rand.json +++ b/tests/fixtures/analysis/registry-matrix/rand.json @@ -24,14 +24,6 @@ "id": "rand#2883ca45911e/ansi_off", "reason": "SELECT rand(INTERVAL '1 02:03:04' DAY TO SECOND) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, - { - "id": "rand#5dd406b51bff/ansi_on", - "reason": "SELECT rand(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "rand#5dd406b51bff/ansi_off", - "reason": "SELECT rand(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, { "id": "rand#a4f856765524/ansi_on", "reason": "SELECT rand(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" diff --git a/tests/fixtures/analysis/registry-matrix/randn.json b/tests/fixtures/analysis/registry-matrix/randn.json index a6b0e817b..d01eca515 100644 --- a/tests/fixtures/analysis/registry-matrix/randn.json +++ b/tests/fixtures/analysis/registry-matrix/randn.json @@ -1,13 +1,5 @@ { "entries": [ - { - "id": "randn#28f8edef6b0c/ansi_on", - "reason": "SELECT randn(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "randn#28f8edef6b0c/ansi_off", - "reason": "SELECT randn(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, { "id": "randn#8df4dfcb761d/ansi_on", "reason": "SELECT randn(map('a', 1)) | KRISHIV_UNSUPPORTED_FEATURE: function map where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" diff --git a/tests/fixtures/analysis/registry-matrix/random.json b/tests/fixtures/analysis/registry-matrix/random.json index bfc049117..83f5dfa94 100644 --- a/tests/fixtures/analysis/registry-matrix/random.json +++ b/tests/fixtures/analysis/registry-matrix/random.json @@ -1,13 +1,5 @@ { "entries": [ - { - "id": "random#56f3c76a60d9/ansi_on", - "reason": "SELECT random(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "random#56f3c76a60d9/ansi_off", - "reason": "SELECT random(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, { "id": "random#6a872774b235/ansi_on", "reason": "SELECT random(INTERVAL '1-2' YEAR TO MONTH) | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" diff --git a/tests/fixtures/analysis/registry-matrix/randstr.json b/tests/fixtures/analysis/registry-matrix/randstr.json index b114ffea2..c38e1ca9d 100644 --- a/tests/fixtures/analysis/registry-matrix/randstr.json +++ b/tests/fixtures/analysis/registry-matrix/randstr.json @@ -8,22 +8,6 @@ "id": "randstr#175c23e922dc/ansi_off", "reason": "SELECT randstr(3, INTERVAL '1-2' YEAR TO MONTH) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, - { - "id": "randstr#380822c97372/ansi_on", - "reason": "SELECT randstr(named_struct('a', 1, 'b', 'x'), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "randstr#380822c97372/ansi_off", - "reason": "SELECT randstr(named_struct('a', 1, 'b', 'x'), 0) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "randstr#42f50df56ea3/ansi_on", - "reason": "SELECT randstr(3, named_struct('a', 1, 'b', 'x')) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "randstr#42f50df56ea3/ansi_off", - "reason": "SELECT randstr(3, named_struct('a', 1, 'b', 'x')) AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, { "id": "randstr#55df7b2e1334/ansi_on", "reason": "SELECT randstr(3, INTERVAL '1 02:03:04' DAY TO SECOND) AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" diff --git a/tests/fixtures/analysis/registry-matrix/typeof.json b/tests/fixtures/analysis/registry-matrix/typeof.json index 71c123e10..57307746d 100644 --- a/tests/fixtures/analysis/registry-matrix/typeof.json +++ b/tests/fixtures/analysis/registry-matrix/typeof.json @@ -24,14 +24,6 @@ "id": "typeof#86964662cdab/ansi_off", "reason": "SELECT typeof(array(1)) | KRISHIV_UNSUPPORTED_FEATURE: function array" }, - { - "id": "typeof#8a4ed8746829/ansi_on", - "reason": "SELECT typeof(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct" - }, - { - "id": "typeof#8a4ed8746829/ansi_off", - "reason": "SELECT typeof(named_struct('a', 1, 'b', 'x')) | KRISHIV_UNSUPPORTED_FEATURE: function named_struct" - }, { "id": "typeof#8b64929c858b/ansi_on", "reason": "SELECT typeof(array(1, 2)) | KRISHIV_UNSUPPORTED_FEATURE: function array" diff --git a/tests/fixtures/analysis/registry-matrix/uniform.json b/tests/fixtures/analysis/registry-matrix/uniform.json index b539b1ee8..05efb925d 100644 --- a/tests/fixtures/analysis/registry-matrix/uniform.json +++ b/tests/fixtures/analysis/registry-matrix/uniform.json @@ -48,14 +48,6 @@ "id": "uniform#454025458e8a/ansi_off", "reason": "SELECT uniform(10, 20, NULL) > 0 AS result | boolean where Spark has boolean not null" }, - { - "id": "uniform#523b3fcbd4d8/ansi_on", - "reason": "SELECT uniform(10, 20, named_struct('a', 1, 'b', 'x')) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "uniform#523b3fcbd4d8/ansi_off", - "reason": "SELECT uniform(10, 20, named_struct('a', 1, 'b', 'x')) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, { "id": "uniform#5656d353ced6/ansi_on", "reason": "SELECT uniform(INTERVAL '1-2' YEAR TO MONTH, 20, 0) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" @@ -88,22 +80,6 @@ "id": "uniform#6afd046a5774/ansi_off", "reason": "SELECT uniform(10, 20, INTERVAL '1-2' YEAR TO MONTH) > 0 AS result | parse: KRISHIV_UNSUPPORTED_FEATURE: INTERVAL literals where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" }, - { - "id": "uniform#8f340576d087/ansi_on", - "reason": "SELECT uniform(10, named_struct('a', 1, 'b', 'x'), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "uniform#8f340576d087/ansi_off", - "reason": "SELECT uniform(10, named_struct('a', 1, 'b', 'x'), 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "uniform#91babb0d6174/ansi_on", - "reason": "SELECT uniform(named_struct('a', 1, 'b', 'x'), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, - { - "id": "uniform#91babb0d6174/ansi_off", - "reason": "SELECT uniform(named_struct('a', 1, 'b', 'x'), 20, 0) > 0 AS result | KRISHIV_UNSUPPORTED_FEATURE: function named_struct where Spark has DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" - }, { "id": "uniform#96fbf5c2c8f2/ansi_on", "reason": "SELECT uniform(10, CAST(1.5 AS DECIMAL(38,18)), 0) > 0 AS result | boolean not null where Spark has boolean" diff --git a/tests/spark-golden/known-failures-cases/order-by.json b/tests/spark-golden/known-failures-cases/order-by.json index 2b7da1b6b..68720aee3 100644 --- a/tests/spark-golden/known-failures-cases/order-by.json +++ b/tests/spark-golden/known-failures-cases/order-by.json @@ -13,11 +13,6 @@ 21, 27 ], - "query-features": [ - 5, - 6, - 12 - ], "variant-type": [ 14 ]