alamb commented on code in PR #8829:
URL: https://github.com/apache/arrow-rs/pull/8829#discussion_r2542135700
##########
arrow-cast/src/pretty.rs:
##########
@@ -1283,4 +1302,473 @@ mod tests {
let actual: Vec<&str> = iso.lines().collect();
assert_eq!(expected_iso, actual, "Actual result:\n{iso}");
}
+
+ //
+ // Custom Formatting
+ //
+
+ /// The factory that will create the [`ArrayFormatter`]s.
+ #[derive(Debug)]
+ struct TestFormatters {}
+
+ impl ArrayFormatterFactory for TestFormatters {
+ fn create_array_formatter<'formatter>(
+ &self,
+ array: &'formatter dyn Array,
+ options: &FormatOptions<'formatter>,
+ field: Option<&'formatter Field>,
+ ) -> Result<Option<ArrayFormatter<'formatter>>, ArrowError> {
+ if field
+ .map(|f| f.extension_type_name() == Some("my_money"))
+ .unwrap_or(false)
+ {
+ // We assume that my_money always is an Int32.
+ let array = array.as_primitive();
+ let display_index = Box::new(MyMoneyFormatter {
+ array,
+ options: options.clone(),
+ });
+ return Ok(Some(ArrayFormatter::new(display_index,
options.safe())));
+ }
+
+ if array.data_type() == &DataType::Int32 {
+ let array = array.as_primitive();
+ let display_index = Box::new(MyInt32Formatter {
+ array,
+ options: options.clone(),
+ });
+ return Ok(Some(ArrayFormatter::new(display_index,
options.safe())));
+ }
+
+ Ok(None)
+ }
+ }
+
+ /// A format that will append a "€" sign to the end of the Int32 values.
+ struct MyMoneyFormatter<'a> {
+ array: &'a Int32Array,
+ options: FormatOptions<'a>,
+ }
+
+ impl<'a> DisplayIndex for MyMoneyFormatter<'a> {
+ fn write(&self, idx: usize, f: &mut dyn Write) ->
crate::display::FormatResult {
+ match self.array.is_valid(idx) {
+ true => write!(f, "{} €", self.array.value(idx))?,
+ false => write!(f, "{}", self.options.null())?,
+ }
+
+ Ok(())
+ }
+ }
+
+ /// The actual formatter
+ struct MyInt32Formatter<'a> {
+ array: &'a Int32Array,
+ options: FormatOptions<'a>,
+ }
+
+ impl<'a> DisplayIndex for MyInt32Formatter<'a> {
+ fn write(&self, idx: usize, f: &mut dyn Write) ->
crate::display::FormatResult {
+ match self.array.is_valid(idx) {
+ true => write!(f, "{} (32-Bit)", self.array.value(idx))?,
+ false => write!(f, "{}", self.options.null())?,
+ }
+
+ Ok(())
+ }
+ }
+
+ #[test]
+ fn test_format_batches_with_custom_formatters() {
+ // define a schema.
+ let options = FormatOptions::new()
+ .with_null("<NULL>")
+ .with_formatter_factory(Some(&TestFormatters {}));
+ let money_metadata = HashMap::from([(
+ extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
+ "my_money".to_owned(),
+ )]);
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("income", DataType::Int32,
true).with_metadata(money_metadata.clone()),
+ ]));
+
+ // define data.
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::new(array::Int32Array::from(vec![
+ Some(1),
+ None,
+ Some(10),
+ Some(100),
+ ]))],
+ )
+ .unwrap();
+
+ let mut buf = String::new();
+ write!(
+ &mut buf,
+ "{}",
+ pretty_format_batches_with_options(&[batch], &options).unwrap()
+ )
+ .unwrap();
+
+ let s = [
+ "+--------+",
+ "| income |",
+ "+--------+",
+ "| 1 € |",
+ "| <NULL> |",
+ "| 10 € |",
+ "| 100 € |",
+ "+--------+",
+ ];
+ let expected = s.join("\n");
+ assert_eq!(expected, buf);
+ }
+
+ #[test]
+ fn test_format_batches_with_custom_formatters_multi_nested_list() {
+ // define a schema.
+ let options = FormatOptions::new()
+ .with_null("<NULL>")
+ .with_formatter_factory(Some(&TestFormatters {}));
+ let money_metadata = HashMap::from([(
+ extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
+ "my_money".to_owned(),
+ )]);
+ let nested_field = Arc::new(
+ Field::new_list_field(DataType::Int32,
true).with_metadata(money_metadata.clone()),
+ );
+
+ // Create nested data
+ let inner_list =
ListBuilder::new(Int32Builder::new()).with_field(nested_field);
+ let mut outer_list = FixedSizeListBuilder::new(inner_list, 2);
+ outer_list.values().append_value([Some(1)]);
+ outer_list.values().append_null();
+ outer_list.append(true);
+ outer_list.values().append_value([Some(2), Some(8)]);
+ outer_list
+ .values()
+ .append_value([Some(50), Some(25), Some(25)]);
+ outer_list.append(true);
+ let outer_list = outer_list.finish();
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "income",
+ outer_list.data_type().clone(),
+ true,
+ )]));
+
+ // define data.
+ let batch = RecordBatch::try_new(schema,
vec![Arc::new(outer_list)]).unwrap();
+
+ let mut buf = String::new();
+ write!(
+ &mut buf,
+ "{}",
+ pretty_format_batches_with_options(&[batch], &options).unwrap()
+ )
+ .unwrap();
+
+ let s = [
+ "+----------------------------------+",
+ "| income |",
+ "+----------------------------------+",
+ "| [[1 €], <NULL>] |",
+ "| [[2 €, 8 €], [50 €, 25 €, 25 €]] |",
+ "+----------------------------------+",
+ ];
+ let expected = s.join("\n");
+ assert_eq!(expected, buf);
+ }
+
+ #[test]
+ fn test_format_batches_with_custom_formatters_nested_struct() {
+ // define a schema.
+ let options = FormatOptions::new()
+ .with_null("<NULL>")
+ .with_formatter_factory(Some(&TestFormatters {}));
+ let money_metadata = HashMap::from([(
+ extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
+ "my_money".to_owned(),
+ )]);
+ let fields = Fields::from(vec![
+ Field::new("name", DataType::Utf8, true),
+ Field::new("income", DataType::Int32,
true).with_metadata(money_metadata.clone()),
+ ]);
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "income",
+ DataType::Struct(fields.clone()),
+ true,
+ )]));
+
+ // Create nested data
+ let mut nested_data = StructBuilder::new(
+ fields,
+ vec![
+ Box::new(StringBuilder::new()),
+ Box::new(Int32Builder::new()),
+ ],
+ );
+ nested_data
+ .field_builder::<StringBuilder>(0)
+ .unwrap()
+ .extend([Some("Gimli"), Some("Legolas"), Some("Aragorn")]);
Review Comment:
⚔️
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
