Support converting RunEndEncodedType to parquet

[brancz]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

I have arrow records that I want to persist as parquet, however, arrow-rs' parquet implementation currently does not support this. When I tried I got this error:

Converting RunEndEncodedType to parquet not supported

Which originates from

arrow-rs/parquet/src/arrow/schema/mod.rs

Lines 777 to 779 in 00a2f73

	DataType::RunEndEncoded(_, _) => Err(arrow_err!(
	"Converting RunEndEncodedType to parquet not supported",
	)),

Describe the solution you'd like

Support converting RunEnd arrays to parquet.

Describe alternatives you've considered

Expanding REE arrays to not be REE encoded, but that kind of defeats the point.

[vegarsti]

Take

[vegarsti]

I would love to make a shot at this.

Do you know what Parquet Type we want as output? Or is this task basically solved once we know that? From what I can tell we want to produce run length encoded Parquet?

Spent some time looking at the code. I need to read up more on Parquet. I couldn't figure out much by staring at the Type struct that is returned by arrow_to_parquet_type returns. Will start by playing with a test with some REE data.

[brancz]

We'd want it to be RLE dictionary encoded (that's what everything is aready defaulting to anyway). I think part of the interesting thing about REE arrays is that it should be relatively straightforward to directly write RLE dictionaries, all we would need to make sure is to cast the values array of the REE array to a dictionary and then encode the data the way parquet expects. It's been a minute since I dealt with encoding parquet data directly, but there should be a fairly straightforward way to convert from one to the other.

[vegarsti]

We'd want it to be RLE dictionary encoded (that's what everything is aready defaulting to anyway). I think part of the interesting thing about REE arrays is that it should be relatively straightforward to directly write RLE dictionaries, all we would need to make sure is to cast the values array of the REE array to a dictionary and then encode the data the way parquet expects. It's been a minute since I dealt with encoding parquet data directly, but there should be a fairly straightforward way to convert from one to the other.

Perfect, thanks, that is helpful.

[vegarsti]

@brancz Do you happen to have a code snippet you could share for trying to persist Arrow as Parquet? That would be very helpful as well.

For the match arm in arrow_to_parquet_type, which is currently panicking, I think doing roughly the same as what the arm above it for the Dictionary type does, will work.

But I haven't found the code for writing records, which I imagine is the more complicated part.

[vegarsti]

@brancz Do you happen to have a code snippet you could share for when you're trying to persist Arrow as Parquet? That would be very helpful as well.

For the match arm in arrow_to_parquet_type, which is currently panicking, I think doing roughly the same as what the arm above it for the Dictionary type does, will work.

But I haven't found the code for converting records, which I imagine is the more complicated part.

I've made a draft PR here just to have a starting point https://github.com/apache/arrow-rs/pull/8069/files#diff-6a684124e78f254fa6ea23b98f04be0a3ff5a2e4b24890d9c6f4dd850ba11333

[vegarsti]

@brancz Do you happen to have a code snippet you could share for trying to persist Arrow as Parquet? That would be very helpful as well.

For the match arm in arrow_to_parquet_type, which is currently panicking, I think doing roughly the same as what the arm above it for the Dictionary type does, will work.

But I haven't found the code for writing records, which I imagine is the more complicated part.

The module I was looking for was parquet's ArrowWriter!

arrow-rs/parquet/src/arrow/arrow_writer/mod.rs

Line 169 in 04f217b

pub struct ArrowWriter<W: Write> {

[vegarsti]

In #8069 now we can write RunArray data to Parquet, but it writes it as plain data, not dictionary data. I basically copied what is done for the Dictionary type where I could find a match clause. I'm guessing what's missing is (at least) some way to transform a RunArray to a DictionaryArray before writing it to Parquet? 🤔

[albertlockett]

We'd like to have this to for otel-arrow's parquet exporter. Happy to help any way I can. I'll help to review #8069 as well.

[albertlockett]

In #8069 now we can write RunArray data to Parquet, but it writes it as plain data, not dictionary data. I basically copied what is done for the Dictionary type where I could find a match clause. I'm guessing what's missing is (at least) some way to transform a RunArray to a DictionaryArray before writing it to Parquet? 🤔

I have a feeling that it might still be using dictionary encoding, but it would be transparent if you're verifying by reading back the parquet file to arrow like is done here:

arrow-rs/parquet/src/arrow/arrow_writer/mod.rs

Line 4332 in 0931768

// Schema of output is plain, not dictionary or REE encoded!!

The reason for this is that, when we create the column encoder, we check if the type supports dictionary encoding here:

arrow-rs/parquet/src/column/writer/encoder.rs

Lines 185 to 187 in 0931768

	let dict_supported = props.dictionary_enabled(descr.path())
	&& has_dictionary_support(T::get_physical_type(), props);
	let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone()));

You can verify this using parquet-tools. For example, if I did something like this

        let run_ends = Int32Array::from_iter_values([5]);
        let all_nulls = UInt32Array::from_iter([None]);
        let run_arr = RunArray::try_new(&run_ends, &all_nulls).unwrap();


        let record_batch = RecordBatch::try_new(
            Arc::new(Schema::new(vec![
                Field::new(
                    "b",
                    DataType::RunEndEncoded(
                        Arc::new(Field::new("run_ends", DataType::Int32, false)),
                        Arc::new(Field::new("values", DataType::UInt32, true)),
                    ),
                    true,   
                )
            ])),
            vec![
                Arc::new(run_arr)
            ]
        ).unwrap();

        let object_store = LocalFileSystem::new_with_prefix("/tmp").unwrap();
        let object_writer = ParquetObjectWriter::new(Arc::new(object_store), "albert.parquet".into());
        let mut arrow_writer = AsyncArrowWriter::try_new_with_options(
            object_writer,
            record_batch.schema(), 
            ArrowWriterOptions::new().with_skip_arrow_metadata(false)
        ).unwrap();

I'd expect to see something like this:

{
  "NumRowGroups": 1,
  "RowGroups": [
    {
      "NumRows": 5,
      "TotalByteSize": 38,
      "Columns": [
        {
          "PathInSchema": [
            "b"
          ],
          "Type": "INT32",
          "Encodings": [
            "PLAIN",
            "RLE",
            "RLE_DICTIONARY"
          ],
          "CompressedSize": 38,
          "UncompressedSize": 38,
          "NumValues": 5,
          "NullCount": 5,
          "CompressionCodec": "UNCOMPRESSED"
        }
      ]
    }
  ]
}

[vegarsti]

Amazing, thank you @albertlockett!