Merge remote-tracking branch 'upstream/main' into broadcast_hash_join

README.md

@@ -113,4 +113,21 @@ INFO src/lib.rs: Comet native library initialized
           +- CometScan parquet [a#14] Batched: true, DataFilters: [isnotnull(a#14), (a#14 > 5)], 
              Format: CometParquet, Location: InMemoryFileIndex(1 paths)[file:/tmp/test], PartitionFilters: [], 
              PushedFilters: [IsNotNull(a), GreaterThan(a,5)], ReadSchema: struct<a:int>
-```
+```
+
+### Enable Comet shuffle
+
+Comet shuffle feature is disabled by default. To enable it, please add related configs:
+
+```
+--conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager
+--conf spark.comet.exec.shuffle.enabled=true
+```
+
+Above configs enable Comet native shuffle which only supports hash partiting and single partition.
+Comet native shuffle doesn't support complext types yet.
+
+To enable columnar shuffle which supports all partitioning and basic complex types, one more config is required:
+```
+--conf spark.comet.columnar.shuffle.enabled=true
+```

common/src/main/java/org/apache/comet/vector/CometPlainVector.java

@@ -105,15 +105,13 @@ public UTF8String getUTF8String(int rowId) {
       int length = Platform.getInt(null, offsetBufferAddress + (rowId + 1L) * 4L) - offset;
       return UTF8String.fromAddress(null, valueBufferAddress + offset, length);
     } else {
-      // Iceberg maps UUID to StringType.
-      // The data type here must be UUID because the only FLBA -> String mapping we have is UUID.
       BaseFixedWidthVector fixedWidthVector = (BaseFixedWidthVector) valueVector;
       int length = fixedWidthVector.getTypeWidth();
       int offset = rowId * length;
       byte[] result = new byte[length];
       Platform.copyMemory(
           null, valueBufferAddress + offset, result, Platform.BYTE_ARRAY_OFFSET, length);
-      return UTF8String.fromString(convertToUuid(result).toString());
+      return UTF8String.fromBytes(result);
     }
   }
 

core/src/common/bit.rs

@@ -131,6 +131,17 @@ pub fn read_num_bytes_u32(size: usize, src: &[u8]) -> u32 {
     trailing_bits(v as u64, size * 8) as u32
 }
 
+/// Similar to the `read_num_bytes` but read nums from bytes in big-endian order
+/// This is used to read bytes from Java's OutputStream which writes bytes in big-endian
+macro_rules! read_num_be_bytes {
+    ($ty:ty, $size:expr, $src:expr) => {{
+        debug_assert!($size <= $src.len());
+        let mut buffer = <$ty as $crate::common::bit::FromBytes>::Buffer::default();
+        buffer.as_mut()[..$size].copy_from_slice(&$src[..$size]);
+        <$ty>::from_be_bytes(buffer)
+    }};
+}
+
 /// Converts value `val` of type `T` to a byte vector, by reading `num_bytes` from `val`.
 /// NOTE: if `val` is less than the size of `T` then it can be truncated.
 #[inline]

core/src/execution/datafusion/expressions/bloom_filter_might_contain.rs

@@ -0,0 +1,152 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::{
+    execution::datafusion::util::spark_bloom_filter::SparkBloomFilter, parquet::data_type::AsBytes,
+};
+use arrow::record_batch::RecordBatch;
+use arrow_array::cast::as_primitive_array;
+use arrow_schema::{DataType, Schema};
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_common::{internal_err, DataFusionError, Result, ScalarValue};
+use datafusion_physical_expr::{aggregate::utils::down_cast_any_ref, PhysicalExpr};
+use std::{
+    any::Any,
+    fmt::Display,
+    hash::{Hash, Hasher},
+    sync::Arc,
+};
+
+/// A physical expression that checks if a value might be in a bloom filter. It corresponds to the
+/// Spark's `BloomFilterMightContain` expression.
+#[derive(Debug, Hash)]
+pub struct BloomFilterMightContain {
+    pub bloom_filter_expr: Arc<dyn PhysicalExpr>,
+    pub value_expr: Arc<dyn PhysicalExpr>,
+    bloom_filter: Option<SparkBloomFilter>,
+}
+
+impl Display for BloomFilterMightContain {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(
+            f,
+            "BloomFilterMightContain [bloom_filter_expr: {}, value_expr: {}]",
+            self.bloom_filter_expr, self.value_expr
+        )
+    }
+}
+
+impl PartialEq<dyn Any> for BloomFilterMightContain {
+    fn eq(&self, _other: &dyn Any) -> bool {
+        down_cast_any_ref(_other)
+            .downcast_ref::<Self>()
+            .map(|other| {
+                self.bloom_filter_expr.eq(&other.bloom_filter_expr)
+                    && self.value_expr.eq(&other.value_expr)
+            })
+            .unwrap_or(false)
+    }
+}
+
+fn evaluate_bloom_filter(
+    bloom_filter_expr: &Arc<dyn PhysicalExpr>,
+) -> Result<Option<SparkBloomFilter>> {
+    // bloom_filter_expr must be a literal/scalar subquery expression, so we can evaluate it
+    // with an empty batch with empty schema
+    let batch = RecordBatch::new_empty(Arc::new(Schema::empty()));
+    let bloom_filter_bytes = bloom_filter_expr.evaluate(&batch)?;
+    match bloom_filter_bytes {
+        ColumnarValue::Scalar(ScalarValue::Binary(v)) => {
+            Ok(v.map(|v| SparkBloomFilter::new(v.as_bytes())))
+        }
+        _ => internal_err!("Bloom filter expression should be evaluated as a scalar binary value"),
+    }
+}
+
+impl BloomFilterMightContain {
+    pub fn try_new(
+        bloom_filter_expr: Arc<dyn PhysicalExpr>,
+        value_expr: Arc<dyn PhysicalExpr>,
+    ) -> Result<Self> {
+        // early evaluate the bloom_filter_expr to get the actual bloom filter
+        let bloom_filter = evaluate_bloom_filter(&bloom_filter_expr)?;
+        Ok(Self {
+            bloom_filter_expr,
+            value_expr,
+            bloom_filter,
+        })
+    }
+}
+
+impl PhysicalExpr for BloomFilterMightContain {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+        Ok(DataType::Boolean)
+    }
+
+    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+        Ok(true)
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+        self.bloom_filter
+            .as_ref()
+            .map(|spark_filter| {
+                let values = self.value_expr.evaluate(batch)?;
+                match values {
+                    ColumnarValue::Array(array) => {
+                        let boolean_array =
+                            spark_filter.might_contain_longs(as_primitive_array(&array));
+                        Ok(ColumnarValue::Array(Arc::new(boolean_array)))
+                    }
+                    ColumnarValue::Scalar(ScalarValue::Int64(v)) => {
+                        let result = v.map(|v| spark_filter.might_contain_long(v));
+                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(result)))
+                    }
+                    _ => internal_err!("value expression should be int64 type"),
+                }
+            })
+            .unwrap_or_else(|| {
+                // when the bloom filter is null, we should return null for all the input
+                Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)))
+            })
+    }
+
+    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+        vec![self.bloom_filter_expr.clone(), self.value_expr.clone()]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        Ok(Arc::new(BloomFilterMightContain::try_new(
+            children[0].clone(),
+            children[1].clone(),
+        )?))
+    }
+
+    fn dyn_hash(&self, state: &mut dyn Hasher) {
+        let mut s = state;
+        self.bloom_filter_expr.hash(&mut s);
+        self.value_expr.hash(&mut s);
+        self.hash(&mut s);
+    }
+}

core/src/execution/datafusion/expressions/mod.rs

@@ -26,6 +26,7 @@ pub mod scalar_funcs;
 pub use normalize_nan::NormalizeNaNAndZero;
 pub mod avg;
 pub mod avg_decimal;
+pub mod bloom_filter_might_contain;
 pub mod strings;
 pub mod subquery;
 pub mod sum_decimal;

core/src/execution/datafusion/mod.rs

@@ -22,3 +22,4 @@ mod operators;
 pub mod planner;
 pub(crate) mod shuffle_writer;
 mod spark_hash;
+mod util;