Merge branch 'main' into experiment_roaring_bitmap_for_int32_anti_semi_joins

Author: coderfender

datafusion/common/src/config.rs

@@ -24,7 +24,7 @@ use crate::encryption::{FileDecryptionProperties, FileEncryptionProperties};
 use crate::error::_config_err;
 use crate::format::{ExplainAnalyzeCategories, ExplainFormat, MetricType};
 use crate::parquet_config::DFParquetWriterVersion;
-use crate::parsers::CompressionTypeVariant;
+use crate::parsers::{CompressionTypeVariant, CsvQuoteStyle};
 use crate::utils::get_available_parallelism;
 use crate::{DataFusionError, Result};
 #[cfg(feature = "parquet_encryption")]
@@ -2042,6 +2042,17 @@ impl ConfigField for CompressionTypeVariant {
     }
 }
 
+impl ConfigField for CsvQuoteStyle {
+    fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
+        v.some(key, self, description)
+    }
+
+    fn set(&mut self, _: &str, value: &str) -> Result<()> {
+        *self = CsvQuoteStyle::from_str(value)?;
+        Ok(())
+    }
+}
+
 /// An implementation trait used to recursively walk configuration
 pub trait Visit {
     fn some<V: Display>(&mut self, key: &str, value: V, description: &'static str);
@@ -3114,6 +3125,15 @@ config_namespace! {
         pub terminator: Option<u8>, default = None
         pub escape: Option<u8>, default = None
         pub double_quote: Option<bool>, default = None
+        /// Quote style for CSV writing.
+        /// One of: "Always", "Necessary", "NonNumeric", "Never"
+        pub quote_style: CsvQuoteStyle, default = CsvQuoteStyle::Necessary
+        /// Whether to ignore leading whitespace in string values when writing CSV.
+        /// Defaults to `false` when `None`.
+        pub ignore_leading_whitespace: Option<bool>, default = None
+        /// Whether to ignore trailing whitespace in string values when writing CSV.
+        /// Defaults to `false` when `None`.
+        pub ignore_trailing_whitespace: Option<bool>, default = None
         /// Specifies whether newlines in (quoted) values are supported.
         ///
         /// Parsing newlines in quoted values may be affected by execution behaviour such as
@@ -3222,6 +3242,30 @@ impl CsvOptions {
         self
     }
 
+    /// Set the quote style for CSV writing.
+    pub fn with_quote_style(mut self, quote_style: CsvQuoteStyle) -> Self {
+        self.quote_style = quote_style;
+        self
+    }
+
+    /// Set whether to ignore leading whitespace in string values when writing CSV.
+    pub fn with_ignore_leading_whitespace(
+        mut self,
+        ignore_leading_whitespace: bool,
+    ) -> Self {
+        self.ignore_leading_whitespace = Some(ignore_leading_whitespace);
+        self
+    }
+
+    /// Set whether to ignore trailing whitespace in string values when writing CSV.
+    pub fn with_ignore_trailing_whitespace(
+        mut self,
+        ignore_trailing_whitespace: bool,
+    ) -> Self {
+        self.ignore_trailing_whitespace = Some(ignore_trailing_whitespace);
+        self
+    }
+
     /// Specifies whether newlines in (quoted) values are supported.
     ///
     /// Parsing newlines in quoted values may be affected by execution behaviour such as

datafusion/common/src/file_options/csv_writer.rs

@@ -94,6 +94,13 @@ impl TryFrom<&CsvOptions> for CsvWriterOptions {
         if let Some(v) = &value.double_quote {
             builder = builder.with_double_quote(*v)
         }
+        builder = builder.with_quote_style(value.quote_style.into());
+        if let Some(v) = &value.ignore_leading_whitespace {
+            builder = builder.with_ignore_leading_whitespace(*v)
+        }
+        if let Some(v) = &value.ignore_trailing_whitespace {
+            builder = builder.with_ignore_trailing_whitespace(*v)
+        }
         Ok(CsvWriterOptions {
             writer_options: builder,
             compression: value.compression,

datafusion/common/src/parsers.rs

@@ -73,3 +73,59 @@ impl CompressionTypeVariant {
         !matches!(self, &Self::UNCOMPRESSED)
     }
 }
+
+/// CSV quote style
+///
+/// Controls when fields are quoted when writing CSV files.
+/// Corresponds to [`arrow::csv::QuoteStyle`].
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
+pub enum CsvQuoteStyle {
+    /// Quote all fields
+    Always,
+    /// Only quote fields when necessary (default)
+    #[default]
+    Necessary,
+    /// Quote all non-numeric fields
+    NonNumeric,
+    /// Never quote fields
+    Never,
+}
+
+impl FromStr for CsvQuoteStyle {
+    type Err = DataFusionError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "always" => Ok(Self::Always),
+            "necessary" => Ok(Self::Necessary),
+            "non_numeric" | "nonnumeric" => Ok(Self::NonNumeric),
+            "never" => Ok(Self::Never),
+            _ => Err(DataFusionError::NotImplemented(format!(
+                "Unsupported CSV quote style {s}"
+            ))),
+        }
+    }
+}
+
+impl From<CsvQuoteStyle> for arrow::csv::QuoteStyle {
+    fn from(style: CsvQuoteStyle) -> Self {
+        match style {
+            CsvQuoteStyle::Always => Self::Always,
+            CsvQuoteStyle::NonNumeric => Self::NonNumeric,
+            CsvQuoteStyle::Never => Self::Never,
+            CsvQuoteStyle::Necessary => Self::Necessary,
+        }
+    }
+}
+
+impl Display for CsvQuoteStyle {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let str = match self {
+            Self::Always => "Always",
+            Self::Necessary => "Necessary",
+            Self::NonNumeric => "NonNumeric",
+            Self::Never => "Never",
+        };
+        write!(f, "{str}")
+    }
+}

datafusion/expr/src/type_coercion/functions.rs

@@ -888,6 +888,8 @@ fn coerced_from<'a>(
         (Utf8View, Utf8 | LargeUtf8 | Null) => Some(type_into.clone()),
         // Any type can be coerced into strings
         (Utf8 | LargeUtf8, _) => Some(type_into.clone()),
+        // We can go into a BinaryView from a Binary or LargeBinary
+        (BinaryView, Binary | LargeBinary | Null) => Some(type_into.clone()),
         (Null, _) if can_cast_types(type_from, type_into) => Some(type_into.clone()),
 
         (List(_), FixedSizeList(_, _)) => Some(type_into.clone()),
@@ -956,6 +958,20 @@ mod tests {
         let cases = vec![
             (DataType::Utf8View, DataType::Utf8),
             (DataType::Utf8View, DataType::LargeUtf8),
+            (DataType::Utf8View, DataType::Null),
+        ];
+
+        for case in cases {
+            assert_eq!(coerced_from(&case.0, &case.1), Some(case.0));
+        }
+    }
+
+    #[test]
+    fn test_binary_conversion() {
+        let cases = vec![
+            (DataType::BinaryView, DataType::Binary),
+            (DataType::BinaryView, DataType::LargeBinary),
+            (DataType::BinaryView, DataType::Null),
         ];
 
         for case in cases {

datafusion/functions-nested/benches/array_remove.rs

@@ -16,12 +16,11 @@
 // under the License.
 
 use arrow::array::{
-    Array, ArrayRef, BinaryArray, BooleanArray, FixedSizeBinaryArray, ListArray,
-    StringArray,
+    Array, ArrayRef, BinaryArray, BooleanArray, Decimal128Array, FixedSizeBinaryArray,
+    Float64Array, Int64Array, ListArray, StringArray,
 };
 use arrow::buffer::OffsetBuffer;
-use arrow::datatypes::{DataType, Decimal128Type, Field, Float64Type, Int64Type};
-use arrow::util::bench_util::create_primitive_list_array_with_seed;
+use arrow::datatypes::{DataType, Field};
 use criterion::{
     criterion_group, criterion_main, {BenchmarkId, Criterion},
 };
@@ -56,15 +55,7 @@ fn bench_array_remove_int64(c: &mut Criterion) {
     let mut group = c.benchmark_group("array_remove_int64");
 
     for &array_size in ARRAY_SIZES {
-        let list_array: ArrayRef =
-            Arc::new(create_primitive_list_array_with_seed::<i32, Int64Type>(
-                NUM_ROWS,
-                0.0,
-                NULL_DENSITY as f32,
-                array_size,
-                SEED,
-            ));
-
+        let list_array = create_int64_list_array(NUM_ROWS, array_size, NULL_DENSITY);
         let element_to_remove = ScalarValue::Int64(Some(1));
         let args = create_args(list_array.clone(), element_to_remove.clone());
 
@@ -105,14 +96,7 @@ fn bench_array_remove_f64(c: &mut Criterion) {
     let mut group = c.benchmark_group("array_remove_f64");
 
     for &array_size in ARRAY_SIZES {
-        let list_array: ArrayRef =
-            Arc::new(create_primitive_list_array_with_seed::<i32, Float64Type>(
-                NUM_ROWS,
-                0.0,
-                NULL_DENSITY as f32,
-                array_size,
-                SEED,
-            ));
+        let list_array = create_f64_list_array(NUM_ROWS, array_size, NULL_DENSITY);
         let element_to_remove = ScalarValue::Float64(Some(1.0));
         let args = create_args(list_array.clone(), element_to_remove.clone());
 
@@ -276,17 +260,8 @@ fn bench_array_remove_decimal64(c: &mut Criterion) {
     let mut group = c.benchmark_group("array_remove_decimal64");
 
     for &array_size in ARRAY_SIZES {
-        let list_array: ArrayRef = Arc::new(create_primitive_list_array_with_seed::<
-            i32,
-            Decimal128Type,
-        >(
-            NUM_ROWS,
-            0.0,
-            NULL_DENSITY as f32,
-            array_size,
-            SEED,
-        ));
-        let element_to_remove = ScalarValue::Decimal128(Some(100_i128), 38, 10);
+        let list_array = create_decimal64_list_array(NUM_ROWS, array_size, NULL_DENSITY);
+        let element_to_remove = ScalarValue::Decimal128(Some(100_i128), 10, 2);
         let args = create_args(list_array.clone(), element_to_remove.clone());
 
         group.bench_with_input(
@@ -301,7 +276,7 @@ fn bench_array_remove_decimal64(c: &mut Criterion) {
                             arg_fields: vec![
                                 Field::new("arr", list_array.data_type().clone(), false)
                                     .into(),
-                                Field::new("el", DataType::Decimal128(38, 10), false)
+                                Field::new("el", DataType::Decimal128(10, 2), false)
                                     .into(),
                             ],
                             number_rows: NUM_ROWS,
@@ -373,6 +348,66 @@ fn create_args(list_array: ArrayRef, element: ScalarValue) -> Vec<ColumnarValue>
     ]
 }
 
+fn create_int64_list_array(
+    num_rows: usize,
+    array_size: usize,
+    null_density: f64,
+) -> ArrayRef {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let values = (0..num_rows * array_size)
+        .map(|_| {
+            if rng.random::<f64>() < null_density {
+                None
+            } else {
+                Some(rng.random_range(0..array_size as i64))
+            }
+        })
+        .collect::<Int64Array>();
+    let offsets = (0..=num_rows)
+        .map(|i| (i * array_size) as i32)
+        .collect::<Vec<i32>>();
+
+    Arc::new(
+        ListArray::try_new(
+            Arc::new(Field::new("item", DataType::Int64, true)),
+            OffsetBuffer::new(offsets.into()),
+            Arc::new(values),
+            None,
+        )
+        .unwrap(),
+    )
+}
+
+fn create_f64_list_array(
+    num_rows: usize,
+    array_size: usize,
+    null_density: f64,
+) -> ArrayRef {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let values = (0..num_rows * array_size)
+        .map(|_| {
+            if rng.random::<f64>() < null_density {
+                None
+            } else {
+                Some(rng.random_range(0..array_size as i64) as f64)
+            }
+        })
+        .collect::<Float64Array>();
+    let offsets = (0..=num_rows)
+        .map(|i| (i * array_size) as i32)
+        .collect::<Vec<i32>>();
+
+    Arc::new(
+        ListArray::try_new(
+            Arc::new(Field::new("item", DataType::Float64, true)),
+            OffsetBuffer::new(offsets.into()),
+            Arc::new(values),
+            None,
+        )
+        .unwrap(),
+    )
+}
+
 fn create_string_list_array(
     num_rows: usize,
     array_size: usize,
@@ -465,6 +500,38 @@ fn create_boolean_list_array(
     )
 }
 
+fn create_decimal64_list_array(
+    num_rows: usize,
+    array_size: usize,
+    null_density: f64,
+) -> ArrayRef {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let values = (0..num_rows * array_size)
+        .map(|_| {
+            if rng.random::<f64>() < null_density {
+                None
+            } else {
+                Some(rng.random_range(0..array_size) as i128 * 100)
+            }
+        })
+        .collect::<Decimal128Array>()
+        .with_precision_and_scale(10, 2)
+        .unwrap();
+    let offsets = (0..=num_rows)
+        .map(|i| (i * array_size) as i32)
+        .collect::<Vec<i32>>();
+
+    Arc::new(
+        ListArray::try_new(
+            Arc::new(Field::new("item", DataType::Decimal128(10, 2), true)),
+            OffsetBuffer::new(offsets.into()),
+            Arc::new(values),
+            None,
+        )
+        .unwrap(),
+    )
+}
+
 fn create_fixed_size_binary_list_array(
     num_rows: usize,
     array_size: usize,