Merge remote-tracking branch 'origin/main' into parquet-reader-bench-baseline

Author: hhhizzz

.github/workflows/rust.yml

@@ -118,7 +118,9 @@ jobs:
         uses: ./.github/actions/setup-builder
       - name: Install cargo-msrv (if needed)
         # cargo-msrv binary may be cached by the cargo cache step in setup-builder, and cargo install will error if it is already installed
-        run: if which cargo-msrv ; then echo "using existing cargo-msrv binary" ; else cargo install cargo-msrv ; fi
+        # --locked uses cargo-msrv's pinned Cargo.lock; without it an unpinned transitive
+        # dep (aws-runtime) resolves to a version that fails to compile (E0282)
+        run: if which cargo-msrv ; then echo "using existing cargo-msrv binary" ; else cargo install cargo-msrv --locked ; fi
       - name: Check all packages
         run: |
           # run `cargo msrv verify --manifest-path "path/to/Cargo.toml"` to see problematic dependencies

Cargo.lock

@@ -34,4 +34,4 @@ crate-type = ["cdylib"]
 
 [dependencies]
 arrow = { path = "../arrow", features = ["pyarrow"] }
-pyo3 = { version = "0.28.0", features = ["extension-module"] }
+pyo3 = { version = "0.29.0", features = ["extension-module"] }

arrow-pyarrow-integration-testing/Cargo.toml

@@ -48,4 +48,4 @@ publish = false
 # Note no dependency on arrow, to ensure arrow-pyarrow can be used by itself
 arrow-array = { path = "../arrow-array" }
 arrow-pyarrow = { path = "../arrow-pyarrow" }
-pyo3 = { version = "0.28.0", default-features = false }
+pyo3 = { version = "0.29.0", default-features = false }

arrow-pyarrow-testing/Cargo.toml

@@ -39,4 +39,4 @@ all-features = true
 arrow-array = { workspace = true, features = ["ffi"] }
 arrow-data = { workspace = true }
 arrow-schema = { workspace = true }
-pyo3 = { version = "0.28.0", default-features = false }
+pyo3 = { version = "0.29.0", default-features = false }

arrow-pyarrow/Cargo.toml

@@ -139,7 +139,9 @@ fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyResult<()> {
         ));
     }
 
-    let capsule_name = unsafe { capsule_name.unwrap().as_cstr().to_str()? };
+    let capsule_name = unsafe { capsule_name.unwrap().as_cstr() }
+        .to_str()
+        .map_err(|e| PyValueError::new_err(e.to_string()))?;
     if capsule_name != name {
         return Err(PyValueError::new_err(format!(
             "Expected name '{name}' in PyCapsule, instead got '{capsule_name}'",

arrow-pyarrow/src/lib.rs

@@ -159,14 +159,14 @@ fn create_string_dictionary_bench_batch(
         true_density,
     )?)
 }
-// commenting out until implementation of RunEndEncoded is complete. See https://github.com/apache/arrow-rs/pull/9936#discussion_r3242936421
-#[allow(dead_code)]
 fn create_ree_bench_batch(
     value_dt: DataType,
     size: usize,
-    null_density: f32,
+    null_pct: Option<u8>,
     true_density: f32,
 ) -> Result<RecordBatch> {
+    const DEFAULT_NULL_PCT: u8 = 10;
+    let null_density = null_pct.unwrap_or(DEFAULT_NULL_PCT) as f32 / 100.0;
     let fields = vec![Field::new(
         "_1",
         DataType::RunEndEncoded(
@@ -506,11 +506,24 @@ fn create_batches() -> Vec<(&'static str, RecordBatch)> {
     let batch = create_string_bench_batch_non_null(BATCH_SIZE, 0.25, 0.75).unwrap();
     batches.push(("string_non_null", batch));
 
-    //let batch = create_ree_bench_batch(DataType::Utf8, BATCH_SIZE, 0.25, 0.75).unwrap();
-    //batches.push(("string_ree", batch));
+    let batch = create_ree_bench_batch(DataType::Utf8, BATCH_SIZE, None, 0.75).unwrap();
+    batches.push(("string_ree", batch));
+
+    let batch = create_ree_bench_batch(DataType::Int32, BATCH_SIZE, None, 0.75).unwrap();
+    batches.push(("int32_ree", batch));
+
+    let batch = create_ree_bench_batch(DataType::Boolean, BATCH_SIZE, None, 0.75).unwrap();
+    batches.push(("bool_ree", batch));
+
+    let batch =
+        create_ree_bench_batch(DataType::FixedSizeBinary(16), BATCH_SIZE, None, 0.75).unwrap();
+    batches.push(("fixed_size_binary_ree", batch));
+
+    let batch = create_ree_bench_batch(DataType::Utf8, BATCH_SIZE, Some(95), 0.75).unwrap();
+    batches.push(("string_ree_95pct_null", batch));
 
-    //let batch = create_ree_bench_batch(DataType::Int32, BATCH_SIZE, 0.25, 0.75).unwrap();
-    //batches.push(("int32_ree", batch));
+    let batch = create_ree_bench_batch(DataType::Int32, BATCH_SIZE, Some(95), 0.75).unwrap();
+    batches.push(("int32_ree_95pct_null", batch));
 
     let batch = create_float_bench_batch_with_nans(BATCH_SIZE, 0.5).unwrap();
     batches.push(("float_with_nans", batch));

parquet/benches/arrow_writer.rs

@@ -44,13 +44,36 @@ use crate::column::chunker::CdcChunk;
 use crate::column::writer::LevelDataRef;
 use crate::errors::{ParquetError, Result};
 use arrow_array::cast::AsArray;
-use arrow_array::{Array, ArrayRef, OffsetSizeTrait};
+use arrow_array::types::RunEndIndexType;
+use arrow_array::{Array, ArrayRef, Int32Array, OffsetSizeTrait, RunArray, downcast_run_array};
 use arrow_buffer::bit_iterator::BitIndexIterator;
 use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow_schema::{DataType, Field};
 use std::ops::Range;
 use std::sync::Arc;
 
+/// Expands a [`DataType::RunEndEncoded`] array into a flat (logical) array of its values type.
+///
+/// use `arrow_select::take` to materialize the  full-length flat array.
+/// This is intentionally simple (O(n)); efficiency can/should be improved
+fn expand_ree_array(array: &ArrayRef) -> Result<ArrayRef> {
+    downcast_run_array!(
+        array => expand_typed_ree(array),
+        _ => unreachable!("expand_ree_array called on non-REE array"),
+    )
+}
+
+fn expand_typed_ree<R: RunEndIndexType>(run_array: &RunArray<R>) -> Result<ArrayRef> {
+    let run_ends = run_array.run_ends();
+    let values = run_array.values();
+    let len = run_array.len();
+    let indices: Int32Array = (0..len)
+        .map(|i| run_ends.get_physical_index(i) as i32)
+        .collect();
+    arrow_select::take::take(values.as_ref(), &indices, None)
+        .map_err(|e| arrow_err!("Failed to expand REE array: {}", e))
+}
+
 /// Performs a depth-first scan of the children of `array`, constructing [`ArrayLevels`]
 /// for each leaf column encountered
 pub(crate) fn calculate_array_levels(array: &ArrayRef, field: &Field) -> Result<Vec<ArrayLevels>> {
@@ -185,6 +208,15 @@ impl LevelInfoBuilder {
                 let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone());
                 Ok(Self::Primitive(levels))
             }
+            DataType::RunEndEncoded(_, value_field) => {
+                let flat = expand_ree_array(array)?;
+                let flat_field = Field::new(
+                    field.name(),
+                    value_field.data_type().clone(),
+                    field.is_nullable(),
+                );
+                Self::try_new(&flat_field, parent_ctx, &flat)
+            }
             DataType::Struct(children) => {
                 let array = array.as_struct();
                 let def_level = match is_nullable {