Merge remote-tracking branch 'apache/main' into fix_uint_cast

Author: alamb

arrow-array/src/array/list_array.rs

@@ -42,16 +42,20 @@ pub trait OffsetSizeTrait: ArrowNativeType + std::ops::AddAssign + Integer {
     const IS_LARGE: bool;
     /// Prefix for the offset size
     const PREFIX: &'static str;
+    /// The max `usize` offset
+    const MAX_OFFSET: usize;
 }
 
 impl OffsetSizeTrait for i32 {
     const IS_LARGE: bool = false;
     const PREFIX: &'static str = "";
+    const MAX_OFFSET: usize = i32::MAX as usize;
 }
 
 impl OffsetSizeTrait for i64 {
     const IS_LARGE: bool = true;
     const PREFIX: &'static str = "Large";
+    const MAX_OFFSET: usize = i64::MAX as usize;
 }
 
 /// An array of [variable length lists], similar to JSON arrays

arrow-array/src/array/struct_array.rs

@@ -91,6 +91,28 @@ impl StructArray {
         Self::try_new(fields, arrays, nulls).unwrap()
     }
 
+    /// Create a new [`StructArray`] from the provided parts, returning an error on failure
+    ///
+    /// The length will be inferred from the length of the child arrays.  Returns an error if
+    /// there are no child arrays.  Consider using [`Self::try_new_with_length`] if the length
+    /// is known to avoid this.
+    ///
+    /// # Errors
+    ///
+    /// Errors if
+    ///
+    /// * `fields.len() == 0`
+    /// * Any reason that [`Self::try_new_with_length`] would error
+    pub fn try_new(
+        fields: Fields,
+        arrays: Vec<ArrayRef>,
+        nulls: Option<NullBuffer>,
+    ) -> Result<Self, ArrowError> {
+        let len = arrays.first().map(|x| x.len()).ok_or_else(||ArrowError::InvalidArgumentError("use StructArray::try_new_with_length or StructArray::new_empty to create a struct array with no fields so that the length can be set correctly".to_string()))?;
+
+        Self::try_new_with_length(fields, arrays, nulls, len)
+    }
+
     /// Create a new [`StructArray`] from the provided parts, returning an error on failure
     ///
     /// # Errors
@@ -102,10 +124,11 @@ impl StructArray {
     /// * `arrays[i].len() != arrays[j].len()`
     /// * `arrays[i].len() != nulls.len()`
     /// * `!fields[i].is_nullable() && !nulls.contains(arrays[i].nulls())`
-    pub fn try_new(
+    pub fn try_new_with_length(
         fields: Fields,
         arrays: Vec<ArrayRef>,
         nulls: Option<NullBuffer>,
+        len: usize,
     ) -> Result<Self, ArrowError> {
         if fields.len() != arrays.len() {
             return Err(ArrowError::InvalidArgumentError(format!(
@@ -114,7 +137,6 @@ impl StructArray {
                 arrays.len()
             )));
         }
-        let len = arrays.first().map(|x| x.len()).unwrap_or_default();
 
         if let Some(n) = nulls.as_ref() {
             if n.len() != len {
@@ -181,6 +203,10 @@ impl StructArray {
 
     /// Create a new [`StructArray`] from the provided parts without validation
     ///
+    /// The length will be inferred from the length of the child arrays.  Panics if there are no
+    /// child arrays.  Consider using [`Self::new_unchecked_with_length`] if the length is known
+    /// to avoid this.
+    ///
     /// # Safety
     ///
     /// Safe if [`Self::new`] would not panic with the given arguments
@@ -193,7 +219,32 @@ impl StructArray {
             return Self::new(fields, arrays, nulls);
         }
 
-        let len = arrays.first().map(|x| x.len()).unwrap_or_default();
+        let len = arrays.first().map(|x| x.len()).expect(
+            "cannot use StructArray::new_unchecked if there are no fields, length is unknown",
+        );
+        Self {
+            len,
+            data_type: DataType::Struct(fields),
+            nulls,
+            fields: arrays,
+        }
+    }
+
+    /// Create a new [`StructArray`] from the provided parts without validation
+    ///
+    /// # Safety
+    ///
+    /// Safe if [`Self::new`] would not panic with the given arguments
+    pub unsafe fn new_unchecked_with_length(
+        fields: Fields,
+        arrays: Vec<ArrayRef>,
+        nulls: Option<NullBuffer>,
+        len: usize,
+    ) -> Self {
+        if cfg!(feature = "force_validate") {
+            return Self::try_new_with_length(fields, arrays, nulls, len).unwrap();
+        }
+
         Self {
             len,
             data_type: DataType::Struct(fields),
@@ -817,9 +868,38 @@ mod tests {
     }
 
     #[test]
+    #[should_panic(expected = "use StructArray::try_new_with_length")]
     fn test_struct_array_from_empty() {
-        let sa = StructArray::from(vec![]);
-        assert!(sa.is_empty())
+        // This can't work because we don't know how many rows the array should have.  Previously we inferred 0 but
+        // that often led to bugs.
+        let _ = StructArray::from(vec![]);
+    }
+
+    #[test]
+    fn test_empty_struct_array() {
+        assert!(StructArray::try_new(Fields::empty(), vec![], None).is_err());
+
+        let arr = StructArray::new_empty_fields(10, None);
+        assert_eq!(arr.len(), 10);
+        assert_eq!(arr.null_count(), 0);
+        assert_eq!(arr.num_columns(), 0);
+
+        let arr2 = StructArray::try_new_with_length(Fields::empty(), vec![], None, 10).unwrap();
+        assert_eq!(arr2.len(), 10);
+
+        let arr = StructArray::new_empty_fields(10, Some(NullBuffer::new_null(10)));
+        assert_eq!(arr.len(), 10);
+        assert_eq!(arr.null_count(), 10);
+        assert_eq!(arr.num_columns(), 0);
+
+        let arr2 = StructArray::try_new_with_length(
+            Fields::empty(),
+            vec![],
+            Some(NullBuffer::new_null(10)),
+            10,
+        )
+        .unwrap();
+        assert_eq!(arr2.len(), 10);
     }
 
     #[test]

arrow-array/src/record_batch.rs

@@ -211,10 +211,11 @@ impl RecordBatch {
     /// Creates a `RecordBatch` from a schema and columns.
     ///
     /// Expects the following:
-    ///  * the vec of columns to not be empty
-    ///  * the schema and column data types to have equal lengths
-    ///    and match
-    ///  * each array in columns to have the same length
+    ///
+    ///  * `!columns.is_empty()`
+    ///  * `schema.fields.len() == columns.len()`
+    ///  * `schema.fields[i].data_type() == columns[i].data_type()`
+    ///  * `columns[i].len() == columns[j].len()`
     ///
     /// If the conditions are not met, an error is returned.
     ///
@@ -240,6 +241,33 @@ impl RecordBatch {
         Self::try_new_impl(schema, columns, &options)
     }
 
+    /// Creates a `RecordBatch` from a schema and columns, without validation.
+    ///
+    /// See [`Self::try_new`] for the checked version.
+    ///
+    /// # Safety
+    ///
+    /// Expects the following:
+    ///
+    ///  * `schema.fields.len() == columns.len()`
+    ///  * `schema.fields[i].data_type() == columns[i].data_type()`
+    ///  * `columns[i].len() == row_count`
+    ///
+    /// Note: if the schema does not match the underlying data exactly, it can lead to undefined
+    /// behavior, for example, via conversion to a `StructArray`, which in turn could lead
+    /// to incorrect access.
+    pub unsafe fn new_unchecked(
+        schema: SchemaRef,
+        columns: Vec<Arc<dyn Array>>,
+        row_count: usize,
+    ) -> Self {
+        Self {
+            schema,
+            columns,
+            row_count,
+        }
+    }
+
     /// Creates a `RecordBatch` from a schema and columns, with additional options,
     /// such as whether to strictly validate field names.
     ///
@@ -340,6 +368,11 @@ impl RecordBatch {
         })
     }
 
+    /// Return the schema, columns and row count of this [`RecordBatch`]
+    pub fn into_parts(self) -> (SchemaRef, Vec<ArrayRef>, usize) {
+        (self.schema, self.columns, self.row_count)
+    }
+
     /// Override the schema of this [`RecordBatch`]
     ///
     /// Returns an error if `schema` is not a superset of the current schema

arrow-avro/src/codec.rs

@@ -51,10 +51,21 @@ impl AvroDataType {
         Field::new(name, d, self.nullability.is_some()).with_metadata(self.metadata.clone())
     }
 
+    /// Returns a reference to the codec used by this data type
+    ///
+    /// The codec determines how Avro data is encoded and mapped to Arrow data types.
+    /// This is useful when we need to inspect or use the specific encoding of a field.
     pub fn codec(&self) -> &Codec {
         &self.codec
     }
 
+    /// Returns the nullability status of this data type
+    ///
+    /// In Avro, nullability is represented through unions with null types.
+    /// The returned value indicates how nulls are encoded in the Avro format:
+    /// - `Some(Nullability::NullFirst)` - Nulls are encoded as the first union variant
+    /// - `Some(Nullability::NullSecond)` - Nulls are encoded as the second union variant
+    /// - `None` - The type is not nullable
     pub fn nullability(&self) -> Option<Nullability> {
         self.nullability
     }
@@ -78,6 +89,10 @@ impl AvroField {
         &self.data_type
     }
 
+    /// Returns the name of this Avro field
+    ///
+    /// This is the field name as defined in the Avro schema.
+    /// It's used to identify fields within a record structure.
     pub fn name(&self) -> &str {
         &self.name
     }
@@ -108,24 +123,46 @@ impl<'a> TryFrom<&Schema<'a>> for AvroField {
 /// <https://avro.apache.org/docs/1.11.1/specification/#encodings>
 #[derive(Debug, Clone)]
 pub enum Codec {
+    /// Represents Avro null type, maps to Arrow's Null data type
     Null,
+    /// Represents Avro boolean type, maps to Arrow's Boolean data type
     Boolean,
+    /// Represents Avro int type, maps to Arrow's Int32 data type
     Int32,
+    /// Represents Avro long type, maps to Arrow's Int64 data type
     Int64,
+    /// Represents Avro float type, maps to Arrow's Float32 data type
     Float32,
+    /// Represents Avro double type, maps to Arrow's Float64 data type
     Float64,
+    /// Represents Avro bytes type, maps to Arrow's Binary data type
     Binary,
+    /// String data represented as UTF-8 encoded bytes, corresponding to Arrow's StringArray
     Utf8,
+    /// Represents Avro date logical type, maps to Arrow's Date32 data type
     Date32,
+    /// Represents Avro time-millis logical type, maps to Arrow's Time32(TimeUnit::Millisecond) data type
     TimeMillis,
+    /// Represents Avro time-micros logical type, maps to Arrow's Time64(TimeUnit::Microsecond) data type
     TimeMicros,
-    /// TimestampMillis(is_utc)
+    /// Represents Avro timestamp-millis or local-timestamp-millis logical type
+    ///
+    /// Maps to Arrow's Timestamp(TimeUnit::Millisecond) data type
+    /// The boolean parameter indicates whether the timestamp has a UTC timezone (true) or is local time (false)
     TimestampMillis(bool),
-    /// TimestampMicros(is_utc)
+    /// Represents Avro timestamp-micros or local-timestamp-micros logical type
+    ///
+    /// Maps to Arrow's Timestamp(TimeUnit::Microsecond) data type
+    /// The boolean parameter indicates whether the timestamp has a UTC timezone (true) or is local time (false)
     TimestampMicros(bool),
+    /// Represents Avro fixed type, maps to Arrow's FixedSizeBinary data type
+    /// The i32 parameter indicates the fixed binary size
     Fixed(i32),
+    /// Represents Avro array type, maps to Arrow's List data type
     List(Arc<AvroDataType>),
+    /// Represents Avro record type, maps to Arrow's Struct data type
     Struct(Arc<[AvroField]>),
+    /// Represents Avro duration logical type, maps to Arrow's Interval(IntervalUnit::MonthDayNano) data type
     Interval,
 }
 

arrow-avro/src/compression.rs

@@ -23,9 +23,16 @@ use std::io::Read;
 pub const CODEC_METADATA_KEY: &str = "avro.codec";
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
+/// Supported compression codecs for Avro data
+///
+/// Avro supports multiple compression formats for data blocks.
+/// This enum represents the compression codecs available in this implementation.
 pub enum CompressionCodec {
+    /// Deflate compression (RFC 1951)
     Deflate,
+    /// Snappy compression
     Snappy,
+    /// ZStandard compression
     ZStandard,
 }
 

arrow-avro/src/lib.rs

@@ -28,11 +28,26 @@
 #![warn(missing_docs)]
 #![allow(unused)] // Temporary
 
+/// Core functionality for reading Avro data into Arrow arrays
+///
+/// Implements the primary reader interface and record decoding logic.
 pub mod reader;
+
+/// Avro schema parsing and representation
+///
+/// Provides types for parsing and representing Avro schema definitions.
 mod schema;
 
+/// Compression codec implementations for Avro
+///
+/// Provides support for various compression algorithms used in Avro files,
+/// including Deflate, Snappy, and ZStandard.
 mod compression;
 
+/// Data type conversions between Avro and Arrow types
+///
+/// This module contains the necessary types and functions to convert between
+/// Avro data types and Arrow data types.
 mod codec;
 
 #[cfg(test)]

arrow-avro/src/reader/record.rs

@@ -37,6 +37,7 @@ pub struct RecordDecoder {
 }
 
 impl RecordDecoder {
+    /// Create a new [`RecordDecoder`] from the provided [`AvroDataType`]
     pub fn try_new(data_type: &AvroDataType) -> Result<Self, ArrowError> {
         match Decoder::try_new(data_type)? {
             Decoder::Record(fields, encodings) => Ok(Self {