feat(functions): uuid, rand, uniform, version, typeof and other misc scalar functions

crates/krishiv-exec/src/context.rs

@@ -398,6 +398,11 @@ pub struct TaskContext {
     name: &'static str,
     memory: OnceLock<MemoryHandle>,
     dependency: Option<GlobalState>,
+    /// 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,
 }
 
 impl fmt::Debug for TaskContext {
@@ -426,9 +431,24 @@ impl TaskContext {
             name,
             memory: OnceLock::new(),
             dependency: None,
+            partition: 0,
         }
     }
 
+    /// 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 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
+    }
+
     pub(crate) fn with_dependency(mut self, state: GlobalState) -> Self {
         self.dependency = Some(state);
         self

crates/krishiv-exec/src/driver.rs

@@ -119,14 +119,16 @@ 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<PipelineRun>) -> TaskHandle {
+/// 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<PipelineRun>, 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 +147,8 @@ pub(crate) fn spawn_driver(run: Arc<PipelineRun>) -> TaskHandle {
     executor.spawn(query.class(), Box::pin(timed))
 }
 
-async fn drive(run: Arc<PipelineRun>) -> Result<()> {
-    Driver::new(run)?.run().await
+async fn drive(run: Arc<PipelineRun>, partition: u32) -> Result<()> {
+    Driver::new(run, partition)?.run().await
 }
 
 /// Times every poll of a driver: the stall detector. A poll is a
@@ -319,7 +321,7 @@ struct Driver {
 }
 
 impl Driver {
-    fn new(run: Arc<PipelineRun>) -> Result<Self> {
+    fn new(run: Arc<PipelineRun>, partition: u32) -> Result<Self> {
         let owner = SpillOwner::new();
         let ctx = |id, name| {
             TaskContext::new(
@@ -329,6 +331,7 @@ impl Driver {
                 id,
                 name,
             )
+            .with_partition(partition)
         };
         let p = &run.pipeline;
         let (source_id, source) = &p.source;

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<ConsumerHandle>,
     query: &QueryContext,
+    partition: u32,
 ) -> Result<Reversed> {
     // 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<ArrayRef> = 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,19 +1680,29 @@ 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");
         assert_eq!(task.probe.rows, 3000);
         // 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());

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<bool>,
+    /// 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<Option<RecordBatch>> {
         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),

crates/krishiv-exec/src/join/probe.rs

@@ -154,6 +154,7 @@ impl Cursor {
         shared: &Shared,
         tables: &Tables<'_>,
         limit: usize,
+        partition: u32,
     ) -> Result<(Option<Batch>, 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<Option<Batch>> {
         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<SpillSignal>,
     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(),
         }))
     }
 

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<LocalState> {
-        Ok(Box::new(Cursor(0)))
+    fn local_state(&self, ctx: &TaskContext) -> Result<LocalState> {
+        Ok(Box::new(State {
+            next: 0,
+            partition: ctx.partition(),
+        }))
     }
 
     fn execute(&self, input: &Batch, local: &mut LocalState) -> Result<OperatorOutput> {
-        let cursor = local_mut::<Cursor>(local, "Expand")?;
+        let state = local_mut::<State>(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)))
     }
 }