Repository: apache/arrow-rs
Pull request #5364 (2024-02-04 13:16:10)

Add example of converting RecordBatches to JSON objects
# Which issue does this PR close?

Related to https://github.com/apache/arrow-rs/pull/5318

# Rationale for this change
 
@tustvold deprecated `record_batches_to_json_rows` in https://github.com/apache/arrow-rs/pull/5318 but there are at least two of us (https://github.com/apache/arrow-rs/pull/5318/files#r1460432887) who are not quite sure how to use the existing APIs or suggested APIs to achieve the same results. 

While converting from arrow --> JSON objects may not be ideal for certain usecases, I think it is a common request so we shouldn't cause users trouble if they were using it

Thus I think adding an example to show that https://github.com/apache/arrow-rs/pull/5318 doesn't regress functionality is warranted 

# What changes are included in this PR?
Add a doc example about how to convert `RecordBatch`es to `serde_json` objects


# Are there any user-facing changes?
Better examples

<!--
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-cast/src/cast.rs b/arrow-cast/src/cast.rs
index bd35096e0..e321ae1b1 100644
--- a/arrow-cast/src/cast.rs
+++ b/arrow-cast/src/cast.rs
@@ -141,6 +141,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
         }
         (_, List(list_to)) => can_cast_types(from_type, list_to.data_type()),
         (_, LargeList(list_to)) => can_cast_types(from_type, list_to.data_type()),
+        (_, FixedSizeList(list_to,size)) if *size == 1 => {
+            can_cast_types(from_type, list_to.data_type())},
         // cast one decimal type to another decimal type
         (Decimal128(_, _), Decimal128(_, _)) => true,
         (Decimal256(_, _), Decimal256(_, _)) => true,
@@ -802,6 +804,9 @@ pub fn cast_with_options(
         }
         (_, List(ref to)) => cast_values_to_list::<i32>(array, to, cast_options),
         (_, LargeList(ref to)) => cast_values_to_list::<i64>(array, to, cast_options),
+        (_, FixedSizeList(ref to, size)) if *size == 1 => {
+            cast_values_to_fixed_size_list(array, to, *size, cast_options)
+        }
         (Decimal128(_, s1), Decimal128(p2, s2)) => {
             cast_decimal_to_decimal_same_type::<Decimal128Type>(
                 array.as_primitive(),
@@ -3040,6 +3045,18 @@ fn cast_values_to_list<O: OffsetSizeTrait>(
     Ok(Arc::new(list))
 }
 
+/// Helper function that takes a primitive array and casts to a fixed size list array.
+fn cast_values_to_fixed_size_list(
+    array: &dyn Array,
+    to: &FieldRef,
+    size: i32,
+    cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError> {
+    let values = cast_with_options(array, to.data_type(), cast_options)?;
+    let list = FixedSizeListArray::new(to.clone(), size, values, None);
+    Ok(Arc::new(list))
+}
+
 /// A specified helper to cast from `GenericBinaryArray` to `GenericStringArray` when they have same
 /// offset size so re-encoding offset is unnecessary.
 fn cast_binary_to_string<O: OffsetSizeTrait>(
@@ -7609,6 +7626,30 @@ mod tests {
         assert_eq!(expected.values(), actual.values());
     }
 
+    #[test]
+    fn test_cast_utf8_to_list() {
+        // DataType::List
+        let array = Arc::new(StringArray::from(vec!["5"])) as ArrayRef;
+        let field = Arc::new(Field::new("", DataType::Int32, false));
+        let list_array = cast(&array, &DataType::List(field.clone())).unwrap();
+        let actual = list_array.as_list_opt::<i32>().unwrap();
+        let expect = ListArray::from_iter_primitive::<Int32Type, _, _>([Some([Some(5)])]);
+        assert_eq!(&expect.value(0), &actual.value(0));
+
+        // DataType::LargeList
+        let list_array = cast(&array, &DataType::LargeList(field.clone())).unwrap();
+        let actual = list_array.as_list_opt::<i64>().unwrap();
+        let expect = LargeListArray::from_iter_primitive::<Int32Type, _, _>([Some([Some(5)])]);
+        assert_eq!(&expect.value(0), &actual.value(0));
+
+        // DataType::FixedSizeList
+        let list_array = cast(&array, &DataType::FixedSizeList(field.clone(), 1)).unwrap();
+        let actual = list_array.as_fixed_size_list_opt().unwrap();
+        let expect =
+            FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>([Some([Some(5)])], 1);
+        assert_eq!(&expect.value(0), &actual.value(0));
+    }
+
     #[test]
     fn test_cast_list_containers() {
         // large-list to list
diff --git a/arrow-json/src/writer.rs b/arrow-json/src/writer.rs
index dd77328cb..acca70290 100644
--- a/arrow-json/src/writer.rs
+++ b/arrow-json/src/writer.rs
@@ -74,7 +74,36 @@
 //! [`LineDelimitedWriter`] and [`ArrayWriter`] will omit writing keys with null values.
 //! In order to explicitly write null values for keys, configure a custom [`Writer`] by
 //! using a [`WriterBuilder`] to construct a [`Writer`].
-
+//!
+//! ## Writing to [serde_json] JSON Objects
+//!
+//! To serialize [`RecordBatch`]es into an array of
+//! [JSON](https://docs.serde.rs/serde_json/) objects you can reparse the resulting JSON string.
+//! Note that this is less efficient than using the `Writer` API.
+//!
+//! ```
+//! # use std::sync::Arc;
+//! # use arrow_array::{Int32Array, RecordBatch};
+//! # use arrow_schema::{DataType, Field, Schema};
+//! let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+//! let a = Int32Array::from(vec![1, 2, 3]);
+//! let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
+//!
+//! // Write the record batch out as json bytes (string)
+//! let buf = Vec::new();
+//! let mut writer = arrow_json::ArrayWriter::new(buf);
+//! writer.write_batches(&vec![&batch]).unwrap();
+//! writer.finish().unwrap();
+//! let json_data = writer.into_inner();
+//!
+//! // Parse the string using serde_json
+//! use serde_json::{Map, Value};
+//! let json_rows: Vec<Map<String, Value>> = serde_json::from_reader(json_data.as_slice()).unwrap();
+//! assert_eq!(
+//!     serde_json::Value::Object(json_rows[1].clone()),
+//!     serde_json::json!({"a": 2}),
+//! );
+//! ```
 mod encoder;
 
 use std::iter;
diff --git a/arrow/src/pyarrow.rs b/arrow/src/pyarrow.rs
index 9a13cfa49..39702ce01 100644
--- a/arrow/src/pyarrow.rs
+++ b/arrow/src/pyarrow.rs
@@ -355,7 +355,16 @@ impl FromPyArrow for RecordBatch {
                 ));
             }
             let array = StructArray::from(array_data);
-            return Ok(array.into());
+            // StructArray does not embed metadata from schema. We need to override
+            // the output schema with the schema from the capsule.
+            let schema = Arc::new(Schema::try_from(schema_ptr).map_err(to_py_err)?);
+            let (_fields, columns, nulls) = array.into_parts();
+            assert_eq!(
+                nulls.map(|n| n.null_count()).unwrap_or_default(),
+                0,
+                "Cannot convert nullable StructArray to RecordBatch, see StructArray documentation"
+            );
+            return RecordBatch::try_new(schema, columns).map_err(to_py_err);
         }
 
         validate_class("RecordBatch", value)?;
diff --git a/parquet/src/file/metadata.rs b/parquet/src/file/metadata.rs
index a1f3c87d0..acd3a9f93 100644
--- a/parquet/src/file/metadata.rs
+++ b/parquet/src/file/metadata.rs
@@ -349,7 +349,13 @@ impl RowGroupMetaData {
 
     /// Method to convert from Thrift.
     pub fn from_thrift(schema_descr: SchemaDescPtr, mut rg: RowGroup) -> Result<RowGroupMetaData> {
-        assert_eq!(schema_descr.num_columns(), rg.columns.len());
+        if schema_descr.num_columns() != rg.columns.len() {
+            return Err(general_err!(
+                "Column count mismatch. Schema has {} columns while Row Group has {}",
+                schema_descr.num_columns(),
+                rg.columns.len()
+            ));
+        }
         let total_byte_size = rg.total_byte_size;
         let num_rows = rg.num_rows;
         let mut columns = vec![];
@@ -1039,6 +1045,75 @@ mod tests {
         }
     }
 
+    /// Test reading a corrupted Parquet file with 3 columns in its schema but only 2 in its row group
+    #[test]
+    fn test_row_group_metadata_thrift_corrupted() {
+        let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
+            SchemaType::group_type_builder("schema")
+                .with_fields(vec![
+                    Arc::new(
+                        SchemaType::primitive_type_builder("a", Type::INT32)
+                            .build()
+                            .unwrap(),
+                    ),
+                    Arc::new(
+                        SchemaType::primitive_type_builder("b", Type::INT32)
+                            .build()
+                            .unwrap(),
+                    ),
+                ])
+                .build()
+                .unwrap(),
+        )));
+
+        let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
+            SchemaType::group_type_builder("schema")
+                .with_fields(vec![
+                    Arc::new(
+                        SchemaType::primitive_type_builder("a", Type::INT32)
+                            .build()
+                            .unwrap(),
+                    ),
+                    Arc::new(
+                        SchemaType::primitive_type_builder("b", Type::INT32)
+                            .build()
+                            .unwrap(),
+                    ),
+                    Arc::new(
+                        SchemaType::primitive_type_builder("c", Type::INT32)
+                            .build()
+                            .unwrap(),
+                    ),
+                ])
+                .build()
+                .unwrap(),
+        )));
+
+        let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
+            .set_num_rows(1000)
+            .set_total_byte_size(2000)
+            .set_column_metadata(vec![
+                ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
+                    .build()
+                    .unwrap(),
+                ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
+                    .build()
+                    .unwrap(),
+            ])
+            .set_ordinal(1)
+            .build()
+            .unwrap();
+
+        let err =
+            RowGroupMetaData::from_thrift(schema_descr_3cols, row_group_meta_2cols.to_thrift())
+                .unwrap_err()
+                .to_string();
+        assert_eq!(
+            err,
+            "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
+        );
+    }
+
     #[test]
     fn test_column_chunk_metadata_thrift_conversion() {
         let column_descr = get_test_schema_descr().column(0);


Test patch:
diff --git a/arrow-pyarrow-integration-testing/tests/test_sql.py b/arrow-pyarrow-integration-testing/tests/test_sql.py
index 16d4e0f12..5320d0a53 100644
--- a/arrow-pyarrow-integration-testing/tests/test_sql.py
+++ b/arrow-pyarrow-integration-testing/tests/test_sql.py
@@ -199,9 +199,10 @@ def test_field_metadata_roundtrip():
 
 def test_schema_roundtrip():
     pyarrow_fields = zip(string.ascii_lowercase, _supported_pyarrow_types)
-    pyarrow_schema = pa.schema(pyarrow_fields)
+    pyarrow_schema = pa.schema(pyarrow_fields, metadata = {b'key1': b'value1'})
     schema = rust.round_trip_schema(pyarrow_schema)
     assert schema == pyarrow_schema
+    assert schema.metadata == pyarrow_schema.metadata
 
 
 def test_primitive_python():
@@ -467,9 +468,11 @@ def test_tensor_array():
     b = rust.round_trip_array(f32_array)
     assert b == f32_array.storage
 
-    batch = pa.record_batch([f32_array], ["tensor"])
+    batch = pa.record_batch([f32_array], ["tensor"], metadata={b'key1': b'value1'})
     b = rust.round_trip_record_batch(batch)
     assert b == batch
+    assert b.schema == batch.schema
+    assert b.schema.metadata == batch.schema.metadata
 
     del b
 
@@ -486,6 +489,7 @@ def test_record_batch_reader():
     b = rust.round_trip_record_batch_reader(a)
 
     assert b.schema == schema
+    assert b.schema.metadata == schema.metadata
     got_batches = list(b)
     assert got_batches == batches
 
@@ -493,6 +497,7 @@ def test_record_batch_reader():
     a = pa.RecordBatchReader.from_batches(schema, batches)
     b = rust.boxed_reader_roundtrip(a)
     assert b.schema == schema
+    assert b.schema.metadata == schema.metadata
     got_batches = list(b)
     assert got_batches == batches
 
@@ -511,6 +516,7 @@ def test_record_batch_reader_pycapsule():
     b = rust.round_trip_record_batch_reader(wrapped)
 
     assert b.schema == schema
+    assert b.schema.metadata == schema.metadata
     got_batches = list(b)
     assert got_batches == batches
 
@@ -519,6 +525,7 @@ def test_record_batch_reader_pycapsule():
     wrapped = StreamWrapper(a)
     b = rust.boxed_reader_roundtrip(wrapped)
     assert b.schema == schema
+    assert b.schema.metadata == schema.metadata
     got_batches = list(b)
     assert got_batches == batches