Skip to content

Commit

Permalink
checkpoint
Browse files Browse the repository at this point in the history
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
  • Loading branch information
jayzhan211 committed Nov 12, 2023
1 parent 89b1f54 commit e27d7f9
Showing 1 changed file with 226 additions and 1 deletion.
227 changes: 226 additions & 1 deletion datafusion/common/src/scalar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1446,7 +1446,55 @@ impl ScalarValue {
let list_array = ScalarValue::iter_to_array_list(scalars)?;
Arc::new(list_array)
}
DataType::Struct(_) => unimplemented!("Not needed for now"),
DataType::Struct(fields) => {
todo!("Struct")
// // Initialize a Vector to store the ScalarValues for each column
// let mut columns: Vec<Vec<ScalarValue>> =
// (0..fields.len()).map(|_| Vec::new()).collect();

// // null mask
// let mut null_mask_builder = BooleanBuilder::new();

// // Iterate over scalars to populate the column scalars for each row
// for scalar in scalars {
// if let ScalarValue::Struct(values, fields) = scalar {
// match values {
// Some(values) => {
// // Push value for each field
// for (column, value) in columns.iter_mut().zip(values) {
// column.push(value.clone());
// }
// null_mask_builder.append_value(false);
// }
// None => {
// // Push NULL of the appropriate type for each field
// for (column, field) in
// columns.iter_mut().zip(fields.as_ref())
// {
// column
// .push(ScalarValue::try_from(field.data_type())?);
// }
// null_mask_builder.append_value(true);
// }
// };
// } else {
// return _internal_err!("Expected Struct but found: {scalar}");
// };
// }

// // Call iter_to_array recursively to convert the scalars for each column into Arrow arrays
// let field_values = fields
// .iter()
// .zip(columns)
// .map(|(field, column)| {
// Ok((field.clone(), Self::iter_to_array(column)?))
// })
// .collect::<Result<Vec<_>>>()?;

// let array = StructArray::from(field_values);
// arrow::compute::nullif(&array, &null_mask_builder.finish())?

}
DataType::Dictionary(key_type, value_type) => {
// create the values array
let value_scalars = scalars
Expand Down Expand Up @@ -4220,6 +4268,183 @@ mod tests {
);
}

// #[test]
// #[cfg(feature = "parquet")]
// fn test_scalar_struct() {
// let field_a = Arc::new(Field::new("A", DataType::Int32, false));
// let field_b = Arc::new(Field::new("B", DataType::Boolean, false));
// let field_c = Arc::new(Field::new("C", DataType::Utf8, false));

// let field_e = Arc::new(Field::new("e", DataType::Int16, false));
// let field_f = Arc::new(Field::new("f", DataType::Int64, false));
// let field_d = Arc::new(Field::new(
// "D",
// DataType::Struct(vec![field_e.clone(), field_f.clone()].into()),
// false,
// ));

// let fields = Fields::from(vec![field_a.clone(), field_b.clone(), field_c.clone(), field_d.clone()]);
// let arrays = vec![
// Arc::new(Int32Array::from(vec![Some(23)])) as ArrayRef,
// Arc::new(BooleanArray::from(vec![Some(false)])) as ArrayRef,
// Arc::new(StringArray::from(vec![Some("Hello")])) as ArrayRef,
// Arc::new(StructArray::from(vec![
// (
// field_e.clone(),
// Arc::new(Int16Array::from(vec![Some(2)])) as ArrayRef,
// ),
// (
// field_f.clone(),
// Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
// ),
// ])) as ArrayRef,
// ];
// let struct_array = StructArray::new(fields, arrays, None);
// let array = Arc::new(struct_array) as ArrayRef;
// let scalar = ScalarValue::Struct(array);

// // Check Display
// assert_eq!(
// format!("{scalar}"),
// String::from(
// "+----------------------------------------------+\n\
// | col |\n\
// +----------------------------------------------+\n\
// | {A: 23, B: false, C: Hello, D: {e: 2, f: 3}} |\n\
// +----------------------------------------------+"
// )
// );

// // Convert to length-2 array
// let array = scalar
// .to_array_of_size(2)
// .expect("Failed to convert to array of size");

// let expected = Arc::new(StructArray::from(vec![
// (
// field_a.clone(),
// Arc::new(Int32Array::from(vec![23, 23])) as ArrayRef,
// ),
// (
// field_b.clone(),
// Arc::new(BooleanArray::from(vec![false, false])) as ArrayRef,
// ),
// (
// field_c.clone(),
// Arc::new(StringArray::from(vec!["Hello", "Hello"])) as ArrayRef,
// ),
// (
// field_d.clone(),
// Arc::new(StructArray::from(vec![
// (
// field_e.clone(),
// Arc::new(Int16Array::from(vec![2, 2])) as ArrayRef,
// ),
// (
// field_f.clone(),
// Arc::new(Int64Array::from(vec![3, 3])) as ArrayRef,
// ),
// ])) as ArrayRef,
// ),
// ])) as ArrayRef;

// assert_eq!(&array, &expected);

// // Construct from second element of ArrayRef
// let constructed = ScalarValue::try_from_array(&expected, 1).unwrap();
// assert_eq!(constructed, scalar);

// // None version
// let none_scalar = ScalarValue::try_from(array.data_type()).unwrap();
// assert!(none_scalar.is_null());

// // Construct with convenience From<Vec<(&str, ScalarValue)>>
// let constructed = ScalarValue::from(vec![
// ("A", ScalarValue::from(23)),
// ("B", ScalarValue::from(false)),
// ("C", ScalarValue::from("Hello")),
// (
// "D",
// ScalarValue::from(vec![
// ("e", ScalarValue::from(2i16)),
// ("f", ScalarValue::from(3i64)),
// ]),
// ),
// ]);
// assert_eq!(constructed, scalar);

// // Build Array from Vec of structs
// let scalars = vec![
// ScalarValue::from(vec![
// ("A", ScalarValue::from(23)),
// ("B", ScalarValue::from(false)),
// ("C", ScalarValue::from("Hello")),
// (
// "D",
// ScalarValue::from(vec![
// ("e", ScalarValue::from(2i16)),
// ("f", ScalarValue::from(3i64)),
// ]),
// ),
// ]),
// ScalarValue::from(vec![
// ("A", ScalarValue::from(7)),
// ("B", ScalarValue::from(true)),
// ("C", ScalarValue::from("World")),
// (
// "D",
// ScalarValue::from(vec![
// ("e", ScalarValue::from(4i16)),
// ("f", ScalarValue::from(5i64)),
// ]),
// ),
// ]),
// ScalarValue::from(vec![
// ("A", ScalarValue::from(-1000)),
// ("B", ScalarValue::from(true)),
// ("C", ScalarValue::from("!!!!!")),
// (
// "D",
// ScalarValue::from(vec![
// ("e", ScalarValue::from(6i16)),
// ("f", ScalarValue::from(7i64)),
// ]),
// ),
// ]),
// ];
// let array = ScalarValue::iter_to_array(scalars).unwrap();

// let expected = Arc::new(StructArray::from(vec![
// (
// field_a,
// Arc::new(Int32Array::from(vec![23, 7, -1000])) as ArrayRef,
// ),
// (
// field_b,
// Arc::new(BooleanArray::from(vec![false, true, true])) as ArrayRef,
// ),
// (
// field_c,
// Arc::new(StringArray::from(vec!["Hello", "World", "!!!!!"])) as ArrayRef,
// ),
// (
// field_d,
// Arc::new(StructArray::from(vec![
// (
// field_e,
// Arc::new(Int16Array::from(vec![2, 4, 6])) as ArrayRef,
// ),
// (
// field_f,
// Arc::new(Int64Array::from(vec![3, 5, 7])) as ArrayRef,
// ),
// ])) as ArrayRef,
// ),
// ])) as ArrayRef;

// assert_eq!(&array, &expected);
// }

#[test]
fn test_lists_in_struct() {
let field_a = Arc::new(Field::new("A", DataType::Utf8, false));
Expand Down

0 comments on commit e27d7f9

Please sign in to comment.