update

Author: luoluoyuyu

src/runtime/streaming/operators/key_operator.rs

@@ -10,13 +10,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//!
+//! Key-by over the physical plan output: key column(s) are **values** projected by the plan
+//! (e.g. `_key_user_id`); **shuffle / `StreamOutput::Keyed` uses `u64` hashes** computed by
+//! [`datafusion_common::hash_utils::create_hashes`] on those columns — same mechanism as
+//! [`crate::runtime::streaming::operators::key_by::KeyByOperator`].
 
 use anyhow::{anyhow, Result};
-use arrow_array::{Array, ArrayRef, RecordBatch, UInt64Array};
+use ahash::RandomState;
 use arrow::compute::{sort_to_indices, take};
+use arrow_array::{Array, ArrayRef, RecordBatch, UInt64Array};
 use async_trait::async_trait;
+use datafusion_common::hash_utils::create_hashes;
+use datafusion_common::ScalarValue;
 use futures::StreamExt;
+use tracing::info;
 
 use crate::runtime::streaming::api::context::TaskContext;
 use crate::runtime::streaming::api::operator::Operator;
@@ -31,6 +38,7 @@ pub struct KeyExecutionOperator {
     name: String,
     executor: StatelessPhysicalExecutor,
     key_fields: Vec<usize>,
+    random_state: RandomState
 }
 
 impl KeyExecutionOperator {
@@ -39,10 +47,18 @@ impl KeyExecutionOperator {
         executor: StatelessPhysicalExecutor,
         key_fields: Vec<usize>,
     ) -> Self {
+        let deterministic_random_state = RandomState::with_seeds(
+            0x1234567890ABCDEF,
+            0x0FEDCBA987654321,
+            0x1357924680135792,
+            0x2468013579246801
+        );
+
         Self {
             name,
             executor,
             key_fields,
+            random_state: deterministic_random_state,
         }
     }
 }
@@ -70,36 +86,28 @@ impl Operator for KeyExecutionOperator {
                 continue;
             }
 
-            let mut final_hashes = vec![0u64; num_rows];
-
-            for &col_idx in &self.key_fields {
-                let col = out_batch.column(col_idx);
-                let int64_array = col
-                    .as_any()
-                    .downcast_ref::<arrow_array::Int64Array>()
-                    .ok_or_else(|| anyhow!("Column at index {} must be Int64Array", col_idx))?;
-
-                for i in 0..num_rows {
-                    let val = int64_array.value(i) as u64;
-                    if self.key_fields.len() == 1 {
-                        final_hashes[i] = val;
-                    } else {
-                        final_hashes[i] ^= val;
-                    }
-                }
-            }
+            let key_arrays: Vec<ArrayRef> = self
+                .key_fields
+                .iter()
+                .map(|&i| out_batch.column(i).clone())
+                .collect();
+
+            let mut hash_buffer = vec![0u64; num_rows];
+            create_hashes(&key_arrays, &self.random_state, &mut hash_buffer)
+                .map_err(|e| anyhow!("KeyExecution failed to hash columns: {e}"))?;
+
+            let hash_array = UInt64Array::from(hash_buffer);
 
-            let hash_array = UInt64Array::from(final_hashes);
             let sorted_indices = sort_to_indices(&hash_array, None, None)
-                .map_err(|e| anyhow!("Failed to sort by key: {e}"))?;
+                .map_err(|e| anyhow!("Failed to sort by hash: {e}"))?;
 
             let sorted_hashes_ref = take(&hash_array, &sorted_indices, None)?;
             let sorted_hashes = sorted_hashes_ref
                 .as_any()
                 .downcast_ref::<UInt64Array>()
                 .unwrap();
 
-            let sorted_columns: std::result::Result<Vec<ArrayRef>, _> = out_batch
+            let sorted_columns: Result<Vec<ArrayRef>, _> = out_batch
                 .columns()
                 .iter()
                 .map(|col| take(col, &sorted_indices, None))
@@ -116,13 +124,11 @@ impl Operator for KeyExecutionOperator {
                 }
 
                 let sub_batch = sorted_batch.slice(start_idx, end_idx - start_idx);
-
                 outputs.push(StreamOutput::Keyed(current_hash, sub_batch));
 
                 start_idx = end_idx;
             }
         }
-
         Ok(outputs)
     }
 
@@ -146,4 +152,3 @@ impl Operator for KeyExecutionOperator {
         Ok(vec![])
     }
 }
-

src/runtime/streaming/operators/windows/session_aggregating_window.rs

@@ -18,9 +18,10 @@ use arrow::compute::{
 use arrow::row::{RowConverter, SortField};
 use arrow_array::types::TimestampNanosecondType;
 use arrow_array::{
-    Array, BooleanArray, PrimitiveArray, RecordBatch, StructArray, TimestampNanosecondArray,
+    Array, ArrayRef, BooleanArray, PrimitiveArray, RecordBatch, StructArray,
+    TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
 };
-use arrow_schema::{DataType, Field, FieldRef, Schema};
+use arrow_schema::{DataType, Field, FieldRef, Schema, TimeUnit};
 use datafusion::execution::context::SessionContext;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::SendableRecordBatchStream;
@@ -30,6 +31,7 @@ use datafusion_proto::protobuf::PhysicalPlanNode;
 use futures::StreamExt;
 use prost::Message;
 use std::collections::{BTreeMap, HashMap, HashSet};
+use tracing::info;
 use std::sync::{Arc, RwLock};
 use std::time::{Duration, SystemTime};
 use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
@@ -315,6 +317,53 @@ fn start_time_for_sorted_batch(batch: &RecordBatch, schema: &FsSchema) -> System
     from_nanos(timestamp_array.value(0) as u128)
 }
 
+/// Appends the stream `_timestamp` column (see [`build_session_output_schema`]) using each
+/// session's `window_end` as the row event time.
+fn append_output_timestamp_column(
+    columns: &mut Vec<ArrayRef>,
+    session_results: &[SessionWindowResult],
+    ts_field: &Field,
+) -> Result<()> {
+    let nanos = |r: &SessionWindowResult| to_nanos(r.window_end)as i64 - 1;
+    match ts_field.data_type() {
+        DataType::Timestamp(TimeUnit::Second, tz) => {
+            let v: Vec<i64> = session_results
+                .iter()
+                .map(|r| (nanos(r) / 1_000_000_000))
+                .collect();
+            columns.push(Arc::new(
+                TimestampSecondArray::from(v).with_timezone_opt(tz.clone()),
+            ));
+        }
+        DataType::Timestamp(TimeUnit::Millisecond, tz) => {
+            let v: Vec<i64> = session_results
+                .iter()
+                .map(|r| (nanos(r) / 1_000_000))
+                .collect();
+            columns.push(Arc::new(
+                TimestampMillisecondArray::from(v).with_timezone_opt(tz.clone()),
+            ));
+        }
+        DataType::Timestamp(TimeUnit::Microsecond, tz) => {
+            let v: Vec<i64> = session_results
+                .iter()
+                .map(|r| (nanos(r) / 1000))
+                .collect();
+            columns.push(Arc::new(
+                TimestampMicrosecondArray::from(v).with_timezone_opt(tz.clone()),
+            ));
+        }
+        DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
+            let v: Vec<i64> = session_results.iter().map(|r| nanos(r)).collect();
+            columns.push(Arc::new(
+                TimestampNanosecondArray::from(v).with_timezone_opt(tz.clone()),
+            ));
+        }
+        dt => bail!("unsupported timestamp type for session window output: {dt}"),
+    }
+    Ok(())
+}
+
 fn build_session_output_schema(
     input: &FsSchema,
     window_field: FieldRef,
@@ -590,6 +639,9 @@ impl SessionWindowOperator {
         columns.insert(self.config.window_index, Arc::new(window_struct_array));
         columns.extend_from_slice(merged_batch.columns());
 
+        let ts_field = self.config.input_schema_ref.schema.field(self.config.input_schema_ref.timestamp_index);
+        append_output_timestamp_column(&mut columns, &session_results, ts_field)?;
+
         RecordBatch::try_new(self.config.output_schema.clone(), columns)
             .context("failed to create session window output batch")
     }