Merge remote-tracking branch 'apache/main' into 7899-avoid-extra-allocation-in-object-builder

Author: alamb

.github/workflows/parquet-variant.yml

@@ -31,6 +31,8 @@ on:
   pull_request:
     paths:
       - parquet-variant/**
+      - parquet-variant-json/**
+      - parquet-variant-compute/**
       - .github/**
 
 jobs:
@@ -50,6 +52,8 @@ jobs:
         run: cargo test -p parquet-variant
       - name: Test parquet-variant-json
         run: cargo test -p parquet-variant-json
+      - name: Test parquet-variant-compute
+        run: cargo test -p parquet-variant-compute
 
   # test compilation
   linux-features:
@@ -63,10 +67,12 @@ jobs:
           submodules: true
       - name: Setup Rust toolchain
         uses: ./.github/actions/setup-builder
-      - name: Check compilation
+      - name: Check compilation (parquet-variant)
         run: cargo check -p parquet-variant
-      - name: Check compilation
+      - name: Check compilation (parquet-variant-json)
         run: cargo check -p parquet-variant-json
+      - name: Check compilation (parquet-variant-compute)
+        run: cargo check -p parquet-variant-compute
 
   clippy:
     name: Clippy
@@ -79,7 +85,9 @@ jobs:
         uses: ./.github/actions/setup-builder
       - name: Setup Clippy
         run: rustup component add clippy
-      - name: Run clippy
+      - name: Run clippy (parquet-variant)
         run: cargo clippy -p parquet-variant --all-targets --all-features -- -D warnings
-      - name: Run clippy
+      - name: Run clippy (parquet-variant-json)
         run: cargo clippy -p parquet-variant-json --all-targets --all-features -- -D warnings
+      - name: Run clippy (parquet-variant-compute)
+        run: cargo clippy -p parquet-variant-compute --all-targets --all-features -- -D warnings

arrow-array/src/builder/mod.rs

@@ -447,6 +447,7 @@ pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<dyn ArrayBuilde
         DataType::Float64 => Box::new(Float64Builder::with_capacity(capacity)),
         DataType::Binary => Box::new(BinaryBuilder::with_capacity(capacity, 1024)),
         DataType::LargeBinary => Box::new(LargeBinaryBuilder::with_capacity(capacity, 1024)),
+        DataType::BinaryView => Box::new(BinaryViewBuilder::with_capacity(capacity)),
         DataType::FixedSizeBinary(len) => {
             Box::new(FixedSizeBinaryBuilder::with_capacity(capacity, *len))
         }
@@ -464,6 +465,7 @@ pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<dyn ArrayBuilde
         ),
         DataType::Utf8 => Box::new(StringBuilder::with_capacity(capacity, 1024)),
         DataType::LargeUtf8 => Box::new(LargeStringBuilder::with_capacity(capacity, 1024)),
+        DataType::Utf8View => Box::new(StringViewBuilder::with_capacity(capacity)),
         DataType::Date32 => Box::new(Date32Builder::with_capacity(capacity)),
         DataType::Date64 => Box::new(Date64Builder::with_capacity(capacity)),
         DataType::Time32(TimeUnit::Second) => {

arrow-ord/src/sort.rs

@@ -180,13 +180,41 @@ where
 
 // partition indices into valid and null indices
 fn partition_validity(array: &dyn Array) -> (Vec<u32>, Vec<u32>) {
-    match array.null_count() {
-        // faster path
-        0 => ((0..(array.len() as u32)).collect(), vec![]),
-        _ => {
-            let indices = 0..(array.len() as u32);
-            indices.partition(|index| array.is_valid(*index as usize))
+    let len = array.len();
+    let null_count = array.null_count();
+    match array.nulls() {
+        Some(nulls) if null_count > 0 => {
+            let mut valid_indices = Vec::with_capacity(len - null_count);
+            let mut null_indices = Vec::with_capacity(null_count);
+
+            let valid_slice = valid_indices.spare_capacity_mut();
+            let null_slice = null_indices.spare_capacity_mut();
+            let mut valid_idx = 0;
+            let mut null_idx = 0;
+
+            nulls.into_iter().enumerate().for_each(|(i, v)| {
+                if v {
+                    valid_slice[valid_idx].write(i as u32);
+                    valid_idx += 1;
+                } else {
+                    null_slice[null_idx].write(i as u32);
+                    null_idx += 1;
+                }
+            });
+
+            assert_eq!(null_idx, null_count);
+            assert_eq!(valid_idx, len - null_count);
+            // Safety: The new lengths match the initial capacity as asserted above,
+            // the bounds checks while writing also ensure they less than or equal to the capacity.
+            unsafe {
+                valid_indices.set_len(valid_idx);
+                null_indices.set_len(null_idx);
+            }
+
+            (valid_indices, null_indices)
         }
+        // faster path
+        _ => ((0..(len as u32)).collect(), vec![]),
     }
 }
 

parquet-variant-compute/Cargo.toml

@@ -41,9 +41,12 @@ name = "parquet_variant_compute"
 bench = false
 
 [dev-dependencies]
+rand = "0.9.1"
 criterion = { version = "0.6", default-features = false }
-rand = { version = "0.9.1" }
+arrow = { workspace = true, features = ["test_utils"] }
+
 
 [[bench]]
-name = "variant_get"
+name = "variant_kernels"
 harness = false
+