Repository: apache/arrow-rs
Pull request #6475 (2024-09-28 02:43:20)

Add take_arrays util for getting entries from 2d arrays
# Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->

Closes #.

# Rationale for this change
 
As in the [suggestion](https://github.com/apache/datafusion/pull/12654#discussion_r1778895596), we think that taking entries from 2d arrays is common enough to move this helper to arrow. 

As an example, if we want to get rows at specific indices for a `RecordBatch`. We can just use
```rust 
take_arrays(batch.columns(), indices, None)
``` 
as opposed to 
```rust
    batch.columns()
        .iter()
        .map(|array| compute::take(array.as_ref(), indices, options.clone()))
        .collect()
```
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->

# What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->

# Are there any user-facing changes?


<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->

<!---
If there are any breaking changes to public APIs, please add the `breaking change` label.
-->



Patch:
diff --git a/arrow-flight/gen/Cargo.toml b/arrow-flight/gen/Cargo.toml
index 08b53b729..aa33f1f49 100644
--- a/arrow-flight/gen/Cargo.toml
+++ b/arrow-flight/gen/Cargo.toml
@@ -34,4 +34,4 @@ publish = false
 # (and checked in) arrow.flight.protocol.rs from changing
 proc-macro2 = { version = "=1.0.86", default-features = false }
 prost-build = { version = "=0.13.3", default-features = false }
-tonic-build = { version = "=0.12.2", default-features = false, features = ["transport", "prost"] }
+tonic-build = { version = "=0.12.3", default-features = false, features = ["transport", "prost"] }
diff --git a/arrow-flight/src/arrow.flight.protocol.rs b/arrow-flight/src/arrow.flight.protocol.rs
index f1eb549d5..0cd4f6948 100644
--- a/arrow-flight/src/arrow.flight.protocol.rs
+++ b/arrow-flight/src/arrow.flight.protocol.rs
@@ -417,7 +417,13 @@ impl CancelStatus {
 }
 /// Generated client implementations.
 pub mod flight_service_client {
-    #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
+    #![allow(
+        unused_variables,
+        dead_code,
+        missing_docs,
+        clippy::wildcard_imports,
+        clippy::let_unit_value,
+    )]
     use tonic::codegen::*;
     use tonic::codegen::http::Uri;
     ///
@@ -521,8 +527,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -555,8 +560,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -590,8 +594,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -640,8 +643,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -672,8 +674,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -704,8 +705,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -736,8 +736,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -767,8 +766,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -801,8 +799,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -832,8 +829,7 @@ pub mod flight_service_client {
                 .ready()
                 .await
                 .map_err(|e| {
-                    tonic::Status::new(
-                        tonic::Code::Unknown,
+                    tonic::Status::unknown(
                         format!("Service was not ready: {}", e.into()),
                     )
                 })?;
@@ -852,7 +848,13 @@ pub mod flight_service_client {
 }
 /// Generated server implementations.
 pub mod flight_service_server {
-    #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
+    #![allow(
+        unused_variables,
+        dead_code,
+        missing_docs,
+        clippy::wildcard_imports,
+        clippy::let_unit_value,
+    )]
     use tonic::codegen::*;
     /// Generated trait containing gRPC methods that should be implemented for use with FlightServiceServer.
     #[async_trait]
@@ -1569,17 +1571,19 @@ pub mod flight_service_server {
                 }
                 _ => {
                     Box::pin(async move {
-                        Ok(
-                            http::Response::builder()
-                                .status(200)
-                                .header("grpc-status", tonic::Code::Unimplemented as i32)
-                                .header(
-                                    http::header::CONTENT_TYPE,
-                                    tonic::metadata::GRPC_CONTENT_TYPE,
-                                )
-                                .body(empty_body())
-                                .unwrap(),
-                        )
+                        let mut response = http::Response::new(empty_body());
+                        let headers = response.headers_mut();
+                        headers
+                            .insert(
+                                tonic::Status::GRPC_STATUS,
+                                (tonic::Code::Unimplemented as i32).into(),
+                            );
+                        headers
+                            .insert(
+                                http::header::CONTENT_TYPE,
+                                tonic::metadata::GRPC_CONTENT_TYPE,
+                            );
+                        Ok(response)
                     })
                 }
             }
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index ed7179fd3..6d037fc41 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -94,6 +94,65 @@ pub fn take(
     }
 }
 
+/// For each [ArrayRef] in the [`Vec<ArrayRef>`], take elements by index and create a new
+/// [`Vec<ArrayRef>`] from those indices.
+///
+/// ```text
+/// ┌────────┬────────┐
+/// │        │        │           ┌────────┐                                ┌────────┬────────┐
+/// │   A    │   1    │           │        │                                │        │        │
+/// ├────────┼────────┤           │   0    │                                │   A    │   1    │
+/// │        │        │           ├────────┤                                ├────────┼────────┤
+/// │   D    │   4    │           │        │                                │        │        │
+/// ├────────┼────────┤           │   2    │  take_arrays(values,indices)   │   B    │   2    │
+/// │        │        │           ├────────┤                                ├────────┼────────┤
+/// │   B    │   2    │           │        │  ───────────────────────────►  │        │        │
+/// ├────────┼────────┤           │   3    │                                │   C    │   3    │
+/// │        │        │           ├────────┤                                ├────────┼────────┤
+/// │   C    │   3    │           │        │                                │        │        │
+/// ├────────┼────────┤           │   1    │                                │   D    │   4    │
+/// │        │        │           └────────┘                                └────────┼────────┘
+/// │   E    │   5    │
+/// └────────┴────────┘
+///    values arrays             indices array                                      result
+/// ```
+///
+/// # Errors
+/// This function errors whenever:
+/// * An index cannot be casted to `usize` (typically 32 bit architectures)
+/// * An index is out of bounds and `options` is set to check bounds.
+///
+/// # Safety
+///
+/// When `options` is not set to check bounds, taking indexes after `len` will panic.
+///
+/// # Examples
+/// ```
+/// # use std::sync::Arc;
+/// # use arrow_array::{StringArray, UInt32Array, cast::AsArray};
+/// # use arrow_select::take::{take, take_arrays};
+/// let string_values = Arc::new(StringArray::from(vec!["zero", "one", "two"]));
+/// let values = Arc::new(UInt32Array::from(vec![0, 1, 2]));
+///
+/// // Take items at index 2, and 1:
+/// let indices = UInt32Array::from(vec![2, 1]);
+/// let taken_arrays = take_arrays(&[string_values, values], &indices, None).unwrap();
+/// let taken_string = taken_arrays[0].as_string::<i32>();
+/// assert_eq!(*taken_string, StringArray::from(vec!["two", "one"]));
+/// let taken_values = taken_arrays[1].as_primitive();
+/// assert_eq!(*taken_values, UInt32Array::from(vec![2, 1]));
+/// ```
+pub fn take_arrays(
+    arrays: &[ArrayRef],
+    indices: &dyn Array,
+    options: Option<TakeOptions>,
+) -> Result<Vec<ArrayRef>, ArrowError> {
+    arrays
+        .iter()
+        .map(|array| take(array.as_ref(), indices, options.clone()))
+        .collect()
+}
+
 /// Verifies that the non-null values of `indices` are all `< len`
 fn check_bounds<T: ArrowPrimitiveType>(
     len: usize,
diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml
index b97b2a571..1d38e67a0 100644
--- a/parquet/Cargo.toml
+++ b/parquet/Cargo.toml
@@ -68,6 +68,7 @@ twox-hash = { version = "1.6", default-features = false }
 paste = { version = "1.0" }
 half = { version = "2.1", default-features = false, features = ["num-traits"] }
 sysinfo = { version = "0.31.2", optional = true, default-features = false, features = ["system"] }
+crc32fast = { version = "1.4.2", optional = true, default-features = false }
 
 [dev-dependencies]
 base64 = { version = "0.22", default-features = false, features = ["std"] }
@@ -117,6 +118,8 @@ object_store = ["dep:object_store", "async"]
 zstd = ["dep:zstd", "zstd-sys"]
 # Display memory in example/write_parquet.rs
 sysinfo = ["dep:sysinfo"]
+# Verify 32-bit CRC checksum when decoding parquet pages
+crc = ["dep:crc32fast"]
 
 [[example]]
 name = "read_parquet"
diff --git a/parquet/README.md b/parquet/README.md
index 0360d15db..a0441ee60 100644
--- a/parquet/README.md
+++ b/parquet/README.md
@@ -60,6 +60,7 @@ The `parquet` crate provides the following features which may be enabled in your
 - `zstd` (default) - support for parquet using `zstd` compression
 - `snap` (default) - support for parquet using `snappy` compression
 - `cli` - parquet [CLI tools](https://github.com/apache/arrow-rs/tree/master/parquet/src/bin)
+- `crc` - enables functionality to automatically verify checksums of each page (if present) when decoding
 - `experimental` - Experimental APIs which may change, even between minor releases
 
 ## Parquet Feature Status
@@ -82,4 +83,4 @@ The `parquet` crate provides the following features which may be enabled in your
 
 ## License
 
-Licensed under the Apache License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0.
+Licensed under the Apache License, Version 2.0: <http://www.apache.org/licenses/LICENSE-2.0>.
diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs
index 6fb0f78c1..b253b73a4 100644
--- a/parquet/src/file/serialized_reader.rs
+++ b/parquet/src/file/serialized_reader.rs
@@ -390,6 +390,15 @@ pub(crate) fn decode_page(
     physical_type: Type,
     decompressor: Option<&mut Box<dyn Codec>>,
 ) -> Result<Page> {
+    // Verify the 32-bit CRC checksum of the page
+    #[cfg(feature = "crc")]
+    if let Some(expected_crc) = page_header.crc {
+        let crc = crc32fast::hash(&buffer);
+        if crc != expected_crc as u32 {
+            return Err(general_err!("Page CRC checksum mismatch"));
+        }
+    }
+
     // When processing data page v2, depending on enabled compression for the
     // page, we should account for uncompressed data ('offset') of
     // repetition and definition levels.


Test patch:
diff --git a/parquet/tests/arrow_reader/checksum.rs b/parquet/tests/arrow_reader/checksum.rs
new file mode 100644
index 000000000..c60908d8b
--- /dev/null
+++ b/parquet/tests/arrow_reader/checksum.rs
@@ -0,0 +1,73 @@
+// 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.
+
+//! This file contains an end to end test for verifying checksums when reading parquet files.
+
+use std::path::PathBuf;
+
+use arrow::util::test_util::parquet_test_data;
+use parquet::arrow::arrow_reader::ArrowReaderBuilder;
+
+#[test]
+fn test_datapage_v1_corrupt_checksum() {
+    let errors = read_file_batch_errors("datapage_v1-corrupt-checksum.parquet");
+    assert_eq!(errors, [
+        Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()), 
+        Ok(()),
+        Ok(()),
+        Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()), 
+        Err("Parquet argument error: Parquet error: Not all children array length are the same!".to_string())
+    ]);
+}
+
+#[test]
+fn test_datapage_v1_uncompressed_checksum() {
+    let errors = read_file_batch_errors("datapage_v1-uncompressed-checksum.parquet");
+    assert_eq!(errors, [Ok(()), Ok(()), Ok(()), Ok(()), Ok(())]);
+}
+
+#[test]
+fn test_datapage_v1_snappy_compressed_checksum() {
+    let errors = read_file_batch_errors("datapage_v1-snappy-compressed-checksum.parquet");
+    assert_eq!(errors, [Ok(()), Ok(()), Ok(()), Ok(()), Ok(())]);
+}
+
+#[test]
+fn test_plain_dict_uncompressed_checksum() {
+    let errors = read_file_batch_errors("plain-dict-uncompressed-checksum.parquet");
+    assert_eq!(errors, [Ok(())]);
+}
+#[test]
+fn test_rle_dict_snappy_checksum() {
+    let errors = read_file_batch_errors("rle-dict-snappy-checksum.parquet");
+    assert_eq!(errors, [Ok(())]);
+}
+
+/// Reads a file and returns a vector with one element per record batch.
+/// The record batch data is replaced with () and errors are stringified.
+fn read_file_batch_errors(name: &str) -> Vec<Result<(), String>> {
+    let path = PathBuf::from(parquet_test_data()).join(name);
+    println!("Reading file: {:?}", path);
+    let file = std::fs::File::open(&path).unwrap();
+    let reader = ArrowReaderBuilder::try_new(file).unwrap().build().unwrap();
+    reader
+        .map(|x| match x {
+            Ok(_) => Ok(()),
+            Err(e) => Err(e.to_string()),
+        })
+        .collect()
+}
diff --git a/parquet/tests/arrow_reader/mod.rs b/parquet/tests/arrow_reader/mod.rs
index cc4c8f3c9..0e6783583 100644
--- a/parquet/tests/arrow_reader/mod.rs
+++ b/parquet/tests/arrow_reader/mod.rs
@@ -36,6 +36,8 @@ use std::sync::Arc;
 use tempfile::NamedTempFile;
 
 mod bad_data;
+#[cfg(feature = "crc")]
+mod checksum;
 mod statistics;
 
 // returns a struct array with columns "int32_col", "float32_col" and "float64_col" with the specified values