// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

//! A comparable row-oriented representation of a collection of [`Array`].

//!

//! [`Row`]s are [normalized for sorting], and can therefore be very efficiently [compared],

//! using [`memcmp`] under the hood, or used in [non-comparison sorts] such as [radix sort].

//! This makes the row format ideal for implementing efficient multi-column sorting,

//! grouping, aggregation, windowing and more, as described in more detail

//! [in this blog post](https://arrow.apache.org/blog/2022/11/07/multi-column-sorts-in-arrow-rust-part-1/).

//!

//! For example, given three input [`Array`], [`RowConverter`] creates byte

//! sequences that [compare] the same as when using [`lexsort`].

//!

//! ```text

//!    ┌─────┐   ┌─────┐   ┌─────┐

//!    │     │   │     │   │     │

//!    ├─────┤ ┌ ┼─────┼ ─ ┼─────┼ ┐              ┏━━━━━━━━━━━━━┓

//!    │     │   │     │   │     │  ─────────────▶┃             ┃

//!    ├─────┤ └ ┼─────┼ ─ ┼─────┼ ┘              ┗━━━━━━━━━━━━━┛

//!    │     │   │     │   │     │

//!    └─────┘   └─────┘   └─────┘

//!                ...

//!    ┌─────┐ ┌ ┬─────┬ ─ ┬─────┬ ┐              ┏━━━━━━━━┓

//!    │     │   │     │   │     │  ─────────────▶┃        ┃

//!    └─────┘ └ ┴─────┴ ─ ┴─────┴ ┘              ┗━━━━━━━━┛

//!     UInt64      Utf8     F64

//!

//!           Input Arrays                          Row Format

//!     (Columns)

//! ```

//!

//! _[`Rows`] must be generated by the same [`RowConverter`] for the comparison

//! to be meaningful._

//!

//! # Basic Example

//! ```

//! # use std::sync::Arc;

//! # use arrow_row::{RowConverter, SortField};

//! # use arrow_array::{ArrayRef, Int32Array, StringArray};

//! # use arrow_array::cast::{AsArray, as_string_array};

//! # use arrow_array::types::Int32Type;

//! # use arrow_schema::DataType;

//!

//! let a1 = Arc::new(Int32Array::from_iter_values([-1, -1, 0, 3, 3])) as ArrayRef;

//! let a2 = Arc::new(StringArray::from_iter_values(["a", "b", "c", "d", "d"])) as ArrayRef;

//! let arrays = vec![a1, a2];

//!

//! // Convert arrays to rows

//! let converter = RowConverter::new(vec![

//!     SortField::new(DataType::Int32),

//!     SortField::new(DataType::Utf8),

//! ]).unwrap();

//! let rows = converter.convert_columns(&arrays).unwrap();

//!

//! // Compare rows

//! for i in 0..4 {

//!     assert!(rows.row(i) <= rows.row(i + 1));

//! }

//! assert_eq!(rows.row(3), rows.row(4));

//!

//! // Convert rows back to arrays

//! let converted = converter.convert_rows(&rows).unwrap();

//! assert_eq!(arrays, converted);

//!

//! // Compare rows from different arrays

//! let a1 = Arc::new(Int32Array::from_iter_values([3, 4])) as ArrayRef;

//! let a2 = Arc::new(StringArray::from_iter_values(["e", "f"])) as ArrayRef;

//! let arrays = vec![a1, a2];

//! let rows2 = converter.convert_columns(&arrays).unwrap();

//!

//! assert!(rows.row(4) < rows2.row(0));

//! assert!(rows.row(4) < rows2.row(1));

//!

//! // Convert selection of rows back to arrays

//! let selection = [rows.row(0), rows2.row(1), rows.row(2), rows2.row(0)];

//! let converted = converter.convert_rows(selection).unwrap();

//! let c1 = converted[0].as_primitive::<Int32Type>();

//! assert_eq!(c1.values(), &[-1, 4, 0, 3]);

//!

//! let c2 = converted[1].as_string::<i32>();

//! let c2_values: Vec<_> = c2.iter().flatten().collect();

//! assert_eq!(&c2_values, &["a", "f", "c", "e"]);

//! ```

//!

//! # Lexicographic Sorts (lexsort)

//!

//! The row format can also be used to implement a fast multi-column / lexicographic sort

//!

//! ```

//! # use arrow_row::{RowConverter, SortField};

//! # use arrow_array::{ArrayRef, UInt32Array};

//! fn lexsort_to_indices(arrays: &[ArrayRef]) -> UInt32Array {

//!     let fields = arrays

//!         .iter()

//!         .map(|a| SortField::new(a.data_type().clone()))

//!         .collect();

//!     let converter = RowConverter::new(fields).unwrap();

//!     let rows = converter.convert_columns(arrays).unwrap();

//!     let mut sort: Vec<_> = rows.iter().enumerate().collect();

//!     sort.sort_unstable_by(|(_, a), (_, b)| a.cmp(b));

//!     UInt32Array::from_iter_values(sort.iter().map(|(i, _)| *i as u32))

//! }

//! ```

//!

//! # Flattening Dictionaries

//!

//! For performance reasons, dictionary arrays are flattened ("hydrated") to their

//! underlying values during row conversion. See [the issue] for more details.

//!

//! This means that the arrays that come out of [`RowConverter::convert_rows`]

//! may not have the same data types as the input arrays. For example, encoding

//! a `Dictionary<Int8, Utf8>` and then will come out as a `Utf8` array.

//!

//! ```

//! # use arrow_array::{Array, ArrayRef, DictionaryArray};

//! # use arrow_array::types::Int8Type;

//! # use arrow_row::{RowConverter, SortField};

//! # use arrow_schema::DataType;

//! # use std::sync::Arc;

//! // Input is a Dictionary array

//! let dict: DictionaryArray::<Int8Type> = ["a", "b", "c", "a", "b"].into_iter().collect();

//! let sort_fields = vec![SortField::new(dict.data_type().clone())];

//! let arrays = vec![Arc::new(dict) as ArrayRef];

//! let converter = RowConverter::new(sort_fields).unwrap();

//! // Convert to rows

//! let rows = converter.convert_columns(&arrays).unwrap();

//! let converted = converter.convert_rows(&rows).unwrap();

//! // result was a Utf8 array, not a Dictionary array

//! assert_eq!(converted[0].data_type(), &DataType::Utf8);

//! ```

//!

//! [non-comparison sorts]: https://en.wikipedia.org/wiki/Sorting_algorithm#Non-comparison_sorts

//! [radix sort]: https://en.wikipedia.org/wiki/Radix_sort

//! [normalized for sorting]: http://wwwlgis.informatik.uni-kl.de/archiv/wwwdvs.informatik.uni-kl.de/courses/DBSREAL/SS2005/Vorlesungsunterlagen/Implementing_Sorting.pdf

//! [`memcmp`]: https://www.man7.org/linux/man-pages/man3/memcmp.3.html

//! [`lexsort`]: https://docs.rs/arrow-ord/latest/arrow_ord/sort/fn.lexsort.html

//! [compared]: PartialOrd

//! [compare]: PartialOrd

//! [the issue]: https://github.com/apache/arrow-rs/issues/4811

#![doc(

    html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg",

    html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg"

)]

#![cfg_attr(docsrs, feature(doc_auto_cfg))]

#![warn(missing_docs)]

use std::cmp::Ordering;

use std::hash::{Hash, Hasher};

use std::sync::Arc;

use arrow_array::cast::*;

use arrow_array::types::ArrowDictionaryKeyType;

use arrow_array::*;

use arrow_buffer::{ArrowNativeType, Buffer, OffsetBuffer, ScalarBuffer};

use arrow_data::{ArrayData, ArrayDataBuilder};

use arrow_schema::*;

use variable::{decode_binary_view, decode_string_view};

use crate::fixed::{decode_bool, decode_fixed_size_binary, decode_primitive};

use crate::list::{compute_lengths_fixed_size_list, encode_fixed_size_list};

use crate::variable::{decode_binary, decode_string};

use arrow_array::types::{Int16Type, Int32Type, Int64Type};

mod fixed;

mod list;

mod run;

mod variable;

/// Converts [`ArrayRef`] columns into a [row-oriented](self) format.

///

/// *Note: The encoding of the row format may change from release to release.*

///

/// ## Overview

///

/// The row format is a variable length byte sequence created by

/// concatenating the encoded form of each column. The encoding for

/// each column depends on its datatype (and sort options).

///

/// The encoding is carefully designed in such a way that escaping is

/// unnecessary: it is never ambiguous as to whether a byte is part of

/// a sentinel (e.g. null) or a value.

///

/// ## Unsigned Integer Encoding

///

/// A null integer is encoded as a `0_u8`, followed by a zero-ed number of bytes corresponding

/// to the integer's length.

///

/// A valid integer is encoded as `1_u8`, followed by the big-endian representation of the

/// integer.

///

/// ```text

///               ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐

///    3          │03│00│00│00│      │01│00│00│00│03│

///               └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘

///               ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐

///   258         │02│01│00│00│      │01│00│00│01│02│

///               └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘

///               ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐

///  23423        │7F│5B│00│00│      │01│00│00│5B│7F│

///               └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘

///               ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐

///  NULL         │??│??│??│??│      │00│00│00│00│00│

///               └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘

///

///              32-bit (4 bytes)        Row Format

///  Value        Little Endian

/// ```

///

/// ## Signed Integer Encoding

///

/// Signed integers have their most significant sign bit flipped, and are then encoded in the

/// same manner as an unsigned integer.

///

/// ```text

///        ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┬──┐

///     5  │05│00│00│00│       │05│00│00│80│       │01│80│00│00│05│

///        └──┴──┴──┴──┘       └──┴──┴──┴──┘       └──┴──┴──┴──┴──┘

///        ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┬──┐

///    -5  │FB│FF│FF│FF│       │FB│FF│FF│7F│       │01│7F│FF│FF│FB│

///        └──┴──┴──┴──┘       └──┴──┴──┴──┘       └──┴──┴──┴──┴──┘

///

///  Value  32-bit (4 bytes)    High bit flipped      Row Format

///          Little Endian

/// ```

///

/// ## Float Encoding

///

/// Floats are converted from IEEE 754 representation to a signed integer representation

/// by flipping all bar the sign bit if they are negative.

///

/// They are then encoded in the same manner as a signed integer.

///

/// ## Fixed Length Bytes Encoding

///

/// Fixed length bytes are encoded in the same fashion as primitive types above.

///

/// For a fixed length array of length `n`:

///

/// A null is encoded as `0_u8` null sentinel followed by `n` `0_u8` bytes

///

/// A valid value is encoded as `1_u8` followed by the value bytes

///

/// ## Variable Length Bytes (including Strings) Encoding

///

/// A null is encoded as a `0_u8`.

///

/// An empty byte array is encoded as `1_u8`.

///

/// A non-null, non-empty byte array is encoded as `2_u8` followed by the byte array

/// encoded using a block based scheme described below.

///

/// The byte array is broken up into fixed-width blocks, each block is written in turn

/// to the output, followed by `0xFF_u8`. The final block is padded to 32-bytes

/// with `0_u8` and written to the output, followed by the un-padded length in bytes

/// of this final block as a `u8`. The first 4 blocks have a length of 8, with subsequent

/// blocks using a length of 32, this is to reduce space amplification for small strings.

///

/// Note the following example encodings use a block size of 4 bytes for brevity:

///

/// ```text

///                       ┌───┬───┬───┬───┬───┬───┐

///  "MEEP"               │02 │'M'│'E'│'E'│'P'│04 │

///                       └───┴───┴───┴───┴───┴───┘

///

///                       ┌───┐

///  ""                   │01 |

///                       └───┘

///

///  NULL                 ┌───┐

///                       │00 │

///                       └───┘

///

/// "Defenestration"      ┌───┬───┬───┬───┬───┬───┐

///                       │02 │'D'│'e'│'f'│'e'│FF │

///                       └───┼───┼───┼───┼───┼───┤

///                           │'n'│'e'│'s'│'t'│FF │

///                           ├───┼───┼───┼───┼───┤

///                           │'r'│'a'│'t'│'r'│FF │

///                           ├───┼───┼───┼───┼───┤

///                           │'a'│'t'│'i'│'o'│FF │

///                           ├───┼───┼───┼───┼───┤

///                           │'n'│00 │00 │00 │01 │

///                           └───┴───┴───┴───┴───┘

/// ```

///

/// This approach is loosely inspired by [COBS] encoding, and chosen over more traditional

/// [byte stuffing] as it is more amenable to vectorisation, in particular AVX-256.

///

/// ## Dictionary Encoding

///

/// Dictionary encoded arrays are hydrated to their underlying values

///

/// ## REE Encoding

///

/// REE (Run End Encoding) arrays, A form of Run Length Encoding, are hydrated to their underlying values.

///

/// ## Struct Encoding

///

/// A null is encoded as a `0_u8`.

///

/// A valid value is encoded as `1_u8` followed by the row encoding of each child.

///

/// This encoding effectively flattens the schema in a depth-first fashion.

///

/// For example

///

/// ```text

/// ┌───────┬────────────────────────┬───────┐

/// │ Int32 │ Struct[Int32, Float32] │ Int32 │

/// └───────┴────────────────────────┴───────┘

/// ```

///

/// Is encoded as

///

/// ```text

/// ┌───────┬───────────────┬───────┬─────────┬───────┐

/// │ Int32 │ Null Sentinel │ Int32 │ Float32 │ Int32 │

/// └───────┴───────────────┴───────┴─────────┴───────┘

/// ```

///

/// ## List Encoding

///

/// Lists are encoded by first encoding all child elements to the row format.

///

/// A list value is then encoded as the concatenation of each of the child elements,

/// separately encoded using the variable length encoding described above, followed

/// by the variable length encoding of an empty byte array.

///

/// For example given:

///

/// ```text

/// [1_u8, 2_u8, 3_u8]

/// [1_u8, null]

/// []

/// null

/// ```

///

/// The elements would be converted to:

///

/// ```text

///     ┌──┬──┐     ┌──┬──┐     ┌──┬──┐     ┌──┬──┐        ┌──┬──┐

///  1  │01│01│  2  │01│02│  3  │01│03│  1  │01│01│  null  │00│00│

///     └──┴──┘     └──┴──┘     └──┴──┘     └──┴──┘        └──┴──┘

///```

///

/// Which would be encoded as

///

/// ```text

///                         ┌──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┐

///  [1_u8, 2_u8, 3_u8]     │02│01│01│00│00│02│02│01│02│00│00│02│02│01│03│00│00│02│01│

///                         └──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┘

///                          └──── 1_u8 ────┘   └──── 2_u8 ────┘  └──── 3_u8 ────┘

///

///                         ┌──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┐

///  [1_u8, null]           │02│01│01│00│00│02│02│00│00│00│00│02│01│

///                         └──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┘

///                          └──── 1_u8 ────┘   └──── null ────┘

///

///```

///

/// With `[]` represented by an empty byte array, and `null` a null byte array.

///

/// ## Fixed Size List Encoding

///

/// Fixed Size Lists are encoded by first encoding all child elements to the row format.

///

/// A non-null list value is then encoded as 0x01 followed by the concatenation of each

/// of the child elements. A null list value is encoded as a null marker.

///

/// For example given:

///

/// ```text

/// [1_u8, 2_u8]

/// [3_u8, null]

/// null

/// ```

///

/// The elements would be converted to:

///

/// ```text

///     ┌──┬──┐     ┌──┬──┐     ┌──┬──┐        ┌──┬──┐

///  1  │01│01│  2  │01│02│  3  │01│03│  null  │00│00│

///     └──┴──┘     └──┴──┘     └──┴──┘        └──┴──┘

///```

///

/// Which would be encoded as

///

/// ```text

///                 ┌──┬──┬──┬──┬──┐

///  [1_u8, 2_u8]   │01│01│01│01│02│

///                 └──┴──┴──┴──┴──┘

///                     └ 1 ┘ └ 2 ┘

///                 ┌──┬──┬──┬──┬──┐

///  [3_u8, null]   │01│01│03│00│00│

///                 └──┴──┴──┴──┴──┘

///                     └ 1 ┘ └null┘

///                 ┌──┐

///  null           │00│

///                 └──┘

///

///```

///

/// # Ordering

///

/// ## Float Ordering

///

/// Floats are totally ordered in accordance to the `totalOrder` predicate as defined

/// in the IEEE 754 (2008 revision) floating point standard.

///

/// The ordering established by this does not always agree with the

/// [`PartialOrd`] and [`PartialEq`] implementations of `f32`. For example,

/// they consider negative and positive zero equal, while this does not

///

/// ## Null Ordering

///

/// The encoding described above will order nulls first, this can be inverted by representing

/// nulls as `0xFF_u8` instead of `0_u8`

///

/// ## Reverse Column Ordering

///

/// The order of a given column can be reversed by negating the encoded bytes of non-null values

///

/// [COBS]: https://en.wikipedia.org/wiki/Consistent_Overhead_Byte_Stuffing

/// [byte stuffing]: https://en.wikipedia.org/wiki/High-Level_Data_Link_Control#Asynchronous_framing

#[derive(Debug)]

pub struct RowConverter {

    fields: Arc<[SortField]>,

    /// State for codecs

    codecs: Vec<Codec>,

}

#[derive(Debug)]

enum Codec {

    /// No additional codec state is necessary

    Stateless,

    /// A row converter for the dictionary values

    /// and the encoding of a row containing only nulls

    Dictionary(RowConverter, OwnedRow),

    /// A row converter for the child fields

    /// and the encoding of a row containing only nulls

    Struct(RowConverter, OwnedRow),

    /// A row converter for the child field

    List(RowConverter),

    /// A row converter for the values array of a run-end encoded array

    RunEndEncoded(RowConverter),

}

impl Codec {

    fn new(sort_field: &SortField) -> Result<Self, ArrowError> {

        match &sort_field.data_type {

            DataType::Dictionary(_, values) => {

                let sort_field =

                    SortField::new_with_options(values.as_ref().clone(), sort_field.options);

                let converter = RowConverter::new(vec![sort_field])?;

                let null_array = new_null_array(values.as_ref(), 1);

                let nulls = converter.convert_columns(&[null_array])?;

                let owned = OwnedRow {

                    data: nulls.buffer.into(),

                    config: nulls.config,

                };

                Ok(Self::Dictionary(converter, owned))

            }

            DataType::RunEndEncoded(_, values) => {

                // Similar to List implementation

                let options = SortOptions {

                    descending: false,

                    nulls_first: sort_field.options.nulls_first != sort_field.options.descending,

                };

                let field = SortField::new_with_options(values.data_type().clone(), options);

                let converter = RowConverter::new(vec![field])?;

                Ok(Self::RunEndEncoded(converter))

            }

            d if !d.is_nested() => Ok(Self::Stateless),

            DataType::List(f) | DataType::LargeList(f) => {

                // The encoded contents will be inverted if descending is set to true

                // As such we set `descending` to false and negate nulls first if it

                // it set to true

                let options = SortOptions {

                    descending: false,

                    nulls_first: sort_field.options.nulls_first != sort_field.options.descending,

                };

                let field = SortField::new_with_options(f.data_type().clone(), options);

                let converter = RowConverter::new(vec![field])?;

                Ok(Self::List(converter))

            }

            DataType::FixedSizeList(f, _) => {

                let field = SortField::new_with_options(f.data_type().clone(), sort_field.options);

                let converter = RowConverter::new(vec![field])?;

                Ok(Self::List(converter))

            }

            DataType::Struct(f) => {

                let sort_fields = f

                    .iter()

                    .map(|x| SortField::new_with_options(x.data_type().clone(), sort_field.options))

                    .collect();

                let converter = RowConverter::new(sort_fields)?;

                let nulls: Vec<_> = f.iter().map(|x| new_null_array(x.data_type(), 1)).collect();

                let nulls = converter.convert_columns(&nulls)?;

                let owned = OwnedRow {

                    data: nulls.buffer.into(),

                    config: nulls.config,

                };

                Ok(Self::Struct(converter, owned))

            }

            _ => Err(ArrowError::NotYetImplemented(format!(

                "not yet implemented: {:?}",

                sort_field.data_type

            ))),

        }

    }

    fn encoder(&self, array: &dyn Array) -> Result<Encoder<'_>, ArrowError> {

        match self {

            Codec::Stateless => Ok(Encoder::Stateless),

            Codec::Dictionary(converter, nulls) => {

                let values = array.as_any_dictionary().values().clone();

                let rows = converter.convert_columns(&[values])?;

                Ok(Encoder::Dictionary(rows, nulls.row()))

            }

            Codec::Struct(converter, null) => {

                let v = as_struct_array(array);

                let rows = converter.convert_columns(v.columns())?;

                Ok(Encoder::Struct(rows, null.row()))

            }

            Codec::List(converter) => {

                let values = match array.data_type() {

                    DataType::List(_) => {

                        let list_array = as_list_array(array);

                        let first_offset = list_array.offsets()[0] as usize;

                        let last_offset =

                            list_array.offsets()[list_array.offsets().len() - 1] as usize;

                        // values can include more data than referenced in the ListArray, only encode

                        // the referenced values.

                        list_array

                            .values()

                            .slice(first_offset, last_offset - first_offset)

                    }

                    DataType::LargeList(_) => {

                        let list_array = as_large_list_array(array);

                        let first_offset = list_array.offsets()[0] as usize;

                        let last_offset =

                            list_array.offsets()[list_array.offsets().len() - 1] as usize;

                        // values can include more data than referenced in the LargeListArray, only encode

                        // the referenced values.

                        list_array

                            .values()

                            .slice(first_offset, last_offset - first_offset)

                    }

                    DataType::FixedSizeList(_, _) => {

                        as_fixed_size_list_array(array).values().clone()

                    }

                    _ => unreachable!(),

                };

                let rows = converter.convert_columns(&[values])?;

                Ok(Encoder::List(rows))

            }

            Codec::RunEndEncoded(converter) => {

                let values = match array.data_type() {

                    DataType::RunEndEncoded(r, _) => match r.data_type() {

                        DataType::Int16 => array.as_run::<Int16Type>().values(),

                        DataType::Int32 => array.as_run::<Int32Type>().values(),

                        DataType::Int64 => array.as_run::<Int64Type>().values(),

                        _ => unreachable!("Unsupported run end index type: {r:?}"),

                    },

                    _ => unreachable!(),

                };

                let rows = converter.convert_columns(std::slice::from_ref(values))?;

                Ok(Encoder::RunEndEncoded(rows))

            }

        }

    }

    fn size(&self) -> usize {

        match self {

            Codec::Stateless => 0,

            Codec::Dictionary(converter, nulls) => converter.size() + nulls.data.len(),

            Codec::Struct(converter, nulls) => converter.size() + nulls.data.len(),

            Codec::List(converter) => converter.size(),

            Codec::RunEndEncoded(converter) => converter.size(),

        }

    }

}

#[derive(Debug)]

enum Encoder<'a> {

    /// No additional encoder state is necessary

    Stateless,

    /// The encoding of the child array and the encoding of a null row

    Dictionary(Rows, Row<'a>),

    /// The row encoding of the child arrays and the encoding of a null row

    ///

    /// It is necessary to encode to a temporary [`Rows`] to avoid serializing

    /// values that are masked by a null in the parent StructArray, otherwise

    /// this would establish an ordering between semantically null values

    Struct(Rows, Row<'a>),

    /// The row encoding of the child array

    List(Rows),

    /// The row encoding of the values array

    RunEndEncoded(Rows),

}

/// Configure the data type and sort order for a given column

#[derive(Debug, Clone, PartialEq, Eq)]

pub struct SortField {

    /// Sort options

    options: SortOptions,

    /// Data type

    data_type: DataType,

}

impl SortField {

    /// Create a new column with the given data type

    pub fn new(data_type: DataType) -> Self {

        Self::new_with_options(data_type, Default::default())

    }

    /// Create a new column with the given data type and [`SortOptions`]

    pub fn new_with_options(data_type: DataType, options: SortOptions) -> Self {

        Self { options, data_type }

    }

    /// Return size of this instance in bytes.

    ///

    /// Includes the size of `Self`.

    pub fn size(&self) -> usize {

        self.data_type.size() + std::mem::size_of::<Self>() - std::mem::size_of::<DataType>()

    }

}

impl RowConverter {

    /// Create a new [`RowConverter`] with the provided schema

    pub fn new(fields: Vec<SortField>) -> Result<Self, ArrowError> {

        if !Self::supports_fields(&fields) {

            return Err(ArrowError::NotYetImplemented(format!(

                "Row format support not yet implemented for: {fields:?}"

            )));

        }

        let codecs = fields.iter().map(Codec::new).collect::<Result<_, _>>()?;

        Ok(Self {

            fields: fields.into(),

            codecs,

        })

    }

    /// Check if the given fields are supported by the row format.

    pub fn supports_fields(fields: &[SortField]) -> bool {

        fields.iter().all(|x| Self::supports_datatype(&x.data_type))

    }

    fn supports_datatype(d: &DataType) -> bool {

        match d {

            _ if !d.is_nested() => true,

            DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => {

                Self::supports_datatype(f.data_type())

            }

            DataType::Struct(f) => f.iter().all(|x| Self::supports_datatype(x.data_type())),

            DataType::RunEndEncoded(_, values) => Self::supports_datatype(values.data_type()),

            _ => false,

        }

    }

    /// Convert [`ArrayRef`] columns into [`Rows`]

    ///

    /// See [`Row`] for information on when [`Row`] can be compared

    ///

    /// See [`Self::convert_rows`] for converting [`Rows`] back into [`ArrayRef`]

    ///

    /// # Panics

    ///

    /// Panics if the schema of `columns` does not match that provided to [`RowConverter::new`]

    pub fn convert_columns(&self, columns: &[ArrayRef]) -> Result<Rows, ArrowError> {

        let num_rows = columns.first().map(|x| x.len()).unwrap_or(0);

        let mut rows = self.empty_rows(num_rows, 0);

        self.append(&mut rows, columns)?;

        Ok(rows)

    }

    /// Convert [`ArrayRef`] columns appending to an existing [`Rows`]

    ///

    /// See [`Row`] for information on when [`Row`] can be compared

    ///

    /// # Panics

    ///

    /// Panics if

    /// * The schema of `columns` does not match that provided to [`RowConverter::new`]

    /// * The provided [`Rows`] were not created by this [`RowConverter`]

    ///

    /// ```

    /// # use std::sync::Arc;

    /// # use std::collections::HashSet;

    /// # use arrow_array::cast::AsArray;

    /// # use arrow_array::StringArray;

    /// # use arrow_row::{Row, RowConverter, SortField};

    /// # use arrow_schema::DataType;

    /// #

    /// let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();

    /// let a1 = StringArray::from(vec!["hello", "world"]);

    /// let a2 = StringArray::from(vec!["a", "a", "hello"]);

    ///

    /// let mut rows = converter.empty_rows(5, 128);

    /// converter.append(&mut rows, &[Arc::new(a1)]).unwrap();

    /// converter.append(&mut rows, &[Arc::new(a2)]).unwrap();

    ///

    /// let back = converter.convert_rows(&rows).unwrap();

    /// let values: Vec<_> = back[0].as_string::<i32>().iter().map(Option::unwrap).collect();

    /// assert_eq!(&values, &["hello", "world", "a", "a", "hello"]);

    /// ```

    pub fn append(&self, rows: &mut Rows, columns: &[ArrayRef]) -> Result<(), ArrowError> {

        assert!(

            Arc::ptr_eq(&rows.config.fields, &self.fields),

            "rows were not produced by this RowConverter"

        );

        if columns.len() != self.fields.len() {

            return Err(ArrowError::InvalidArgumentError(format!(

                "Incorrect number of arrays provided to RowConverter, expected {} got {}",

                self.fields.len(),

                columns.len()

            )));

        }

        for colum in columns.iter().skip(1) {

            if colum.len() != columns[0].len() {

                return Err(ArrowError::InvalidArgumentError(format!(

                    "RowConverter columns must all have the same length, expected {} got {}",

                    columns[0].len(),

                    colum.len()

                )));

            }

        }

        let encoders = columns

            .iter()

            .zip(&self.codecs)

            .zip(self.fields.iter())

            .map(|((column, codec), field)| {

                if !column.data_type().equals_datatype(&field.data_type) {

                    return Err(ArrowError::InvalidArgumentError(format!(

                        "RowConverter column schema mismatch, expected {} got {}",

                        field.data_type,

                        column.data_type()

                    )));

                }

                codec.encoder(column.as_ref())

            })

            .collect::<Result<Vec<_>, _>>()?;

        let write_offset = rows.num_rows();

        let lengths = row_lengths(columns, &encoders);

        let total = lengths.extend_offsets(rows.offsets[write_offset], &mut rows.offsets);

        rows.buffer.resize(total, 0);

        for ((column, field), encoder) in columns.iter().zip(self.fields.iter()).zip(encoders) {

            // We encode a column at a time to minimise dispatch overheads

            encode_column(

                &mut rows.buffer,

                &mut rows.offsets[write_offset..],

                column.as_ref(),

                field.options,

                &encoder,

            )

        }

        if cfg!(debug_assertions) {

            assert_eq!(*rows.offsets.last().unwrap(), rows.buffer.len());

            rows.offsets

                .windows(2)

                .for_each(|w| assert!(w[0] <= w[1], "offsets should be monotonic"));

        }

        Ok(())

    }

    /// Convert [`Rows`] columns into [`ArrayRef`]

    ///

    /// See [`Self::convert_columns`] for converting [`ArrayRef`] into [`Rows`]

    ///

    /// # Panics

    ///

    /// Panics if the rows were not produced by this [`RowConverter`]

    pub fn convert_rows<'a, I>(&self, rows: I) -> Result<Vec<ArrayRef>, ArrowError>

    where

        I: IntoIterator<Item = Row<'a>>,

    {

        let mut validate_utf8 = false;

        let mut rows: Vec<_> = rows

            .into_iter()

            .map(|row| {

                assert!(

                    Arc::ptr_eq(&row.config.fields, &self.fields),

                    "rows were not produced by this RowConverter"

                );

                validate_utf8 |= row.config.validate_utf8;

                row.data

            })

            .collect();

        // SAFETY

        // We have validated that the rows came from this [`RowConverter`]

        // and therefore must be valid

        let result = unsafe { self.convert_raw(&mut rows, validate_utf8) }?;

        if cfg!(test) {

            for (i, row) in rows.iter().enumerate() {

                if !row.is_empty() {

                    return Err(ArrowError::InvalidArgumentError(format!(

                        "Codecs {codecs:?} did not consume all bytes for row {i}, remaining bytes: {row:?}",

                        codecs = &self.codecs

                    )));

                }

            }

        }

        Ok(result)

    }

    /// Returns an empty [`Rows`] with capacity for `row_capacity` rows with

    /// a total length of `data_capacity`

    ///

    /// This can be used to buffer a selection of [`Row`]

    ///

    /// ```

    /// # use std::sync::Arc;

    /// # use std::collections::HashSet;

    /// # use arrow_array::cast::AsArray;

    /// # use arrow_array::StringArray;

    /// # use arrow_row::{Row, RowConverter, SortField};

    /// # use arrow_schema::DataType;

    /// #

    /// let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();

    /// let array = StringArray::from(vec!["hello", "world", "a", "a", "hello"]);

    ///

    /// // Convert to row format and deduplicate

    /// let converted = converter.convert_columns(&[Arc::new(array)]).unwrap();

    /// let mut distinct_rows = converter.empty_rows(3, 100);

    /// let mut dedup: HashSet<Row> = HashSet::with_capacity(3);

    /// converted.iter().filter(|row| dedup.insert(*row)).for_each(|row| distinct_rows.push(row));

    ///

    /// // Note: we could skip buffering and feed the filtered iterator directly

    /// // into convert_rows, this is done for demonstration purposes only

    /// let distinct = converter.convert_rows(&distinct_rows).unwrap();

    /// let values: Vec<_> = distinct[0].as_string::<i32>().iter().map(Option::unwrap).collect();

    /// assert_eq!(&values, &["hello", "world", "a"]);

    /// ```

    pub fn empty_rows(&self, row_capacity: usize, data_capacity: usize) -> Rows {

        let mut offsets = Vec::with_capacity(row_capacity.saturating_add(1));

        offsets.push(0);

        Rows {

            offsets,

            buffer: Vec::with_capacity(data_capacity),

            config: RowConfig {

                fields: self.fields.clone(),

                validate_utf8: false,

            },

        }

    }

    /// Create a new [Rows] instance from the given binary data.

    ///

    /// ```

    /// # use std::sync::Arc;

    /// # use std::collections::HashSet;

    /// # use arrow_array::cast::AsArray;

    /// # use arrow_array::StringArray;

    /// # use arrow_row::{OwnedRow, Row, RowConverter, RowParser, SortField};

    /// # use arrow_schema::DataType;

    /// #

    /// let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();

    /// let array = StringArray::from(vec!["hello", "world", "a", "a", "hello"]);

    /// let rows = converter.convert_columns(&[Arc::new(array)]).unwrap();

    ///

    /// // We can convert rows into binary format and back in batch.

    /// let values: Vec<OwnedRow> = rows.iter().map(|r| r.owned()).collect();

    /// let binary = rows.try_into_binary().expect("known-small array");

    /// let converted = converter.from_binary(binary.clone());

    /// assert!(converted.iter().eq(values.iter().map(|r| r.row())));

    /// ```

    ///

    /// # Panics

    ///

    /// This function expects the passed [BinaryArray] to contain valid row data as produced by this

    /// [RowConverter]. It will panic if any rows are null. Operations on the returned [Rows] may

    /// panic if the data is malformed.

    pub fn from_binary(&self, array: BinaryArray) -> Rows {

        assert_eq!(

            array.null_count(),

            0,

            "can't construct Rows instance from array with nulls"

        );

        Rows {

            buffer: array.values().to_vec(),

            offsets: array.offsets().iter().map(|&i| i.as_usize()).collect(),

            config: RowConfig {

                fields: Arc::clone(&self.fields),

                validate_utf8: true,

            },

        }

    }

    /// Convert raw bytes into [`ArrayRef`]

    ///

    /// # Safety

    ///

    /// `rows` must contain valid data for this [`RowConverter`]

    unsafe fn convert_raw(

        &self,

        rows: &mut [&[u8]],

        validate_utf8: bool,

    ) -> Result<Vec<ArrayRef>, ArrowError> {

        self.fields

            .iter()

            .zip(&self.codecs)

            .map(|(field, codec)| decode_column(field, rows, codec, validate_utf8))

            .collect()

    }

    /// Returns a [`RowParser`] that can be used to parse [`Row`] from bytes

    pub fn parser(&self) -> RowParser {

        RowParser::new(Arc::clone(&self.fields))

    }

    /// Returns the size of this instance in bytes

    ///

    /// Includes the size of `Self`.

    pub fn size(&self) -> usize {

        std::mem::size_of::<Self>()

            + self.fields.iter().map(|x| x.size()).sum::<usize>()

            + self.codecs.capacity() * std::mem::size_of::<Codec>()

            + self.codecs.iter().map(Codec::size).sum::<usize>()

    }

}

/// A [`RowParser`] can be created from a [`RowConverter`] and used to parse bytes to [`Row`]

#[derive(Debug)]

pub struct RowParser {

    config: RowConfig,

}

impl RowParser {

    fn new(fields: Arc<[SortField]>) -> Self {

        Self {

            config: RowConfig {

                fields,

                validate_utf8: true,

            },

        }

    }

    /// Creates a [`Row`] from the provided `bytes`.

    ///

    /// `bytes` must be a [`Row`] produced by the [`RowConverter`] associated with

    /// this [`RowParser`], otherwise subsequent operations with the produced [`Row`] may panic

    pub fn parse<'a>(&'a self, bytes: &'a [u8]) -> Row<'a> {

        Row {

            data: bytes,

            config: &self.config,

        }

    }

}

/// The config of a given set of [`Row`]

#[derive(Debug, Clone)]

struct RowConfig {

    /// The schema for these rows

    fields: Arc<[SortField]>,

    /// Whether to run UTF-8 validation when converting to arrow arrays

    validate_utf8: bool,

}

/// A row-oriented representation of arrow data, that is normalized for comparison.

///

/// See the [module level documentation](self) and [`RowConverter`] for more details.

#[derive(Debug)]

pub struct Rows {

    /// Underlying row bytes

    buffer: Vec<u8>,

    /// Row `i` has data `&buffer[offsets[i]..offsets[i+1]]`

    offsets: Vec<usize>,

    /// The config for these rows

    config: RowConfig,

}

impl Rows {

    /// Append a [`Row`] to this [`Rows`]

    pub fn push(&mut self, row: Row<'_>) {

        assert!(

            Arc::ptr_eq(&row.config.fields, &self.config.fields),

            "row was not produced by this RowConverter"

        );

        self.config.validate_utf8 |= row.config.validate_utf8;

        self.buffer.extend_from_slice(row.data);

        self.offsets.push(self.buffer.len())

    }

    /// Returns the row at index `row`

    pub fn row(&self, row: usize) -> Row<'_> {

        assert!(row + 1 < self.offsets.len());

        unsafe { self.row_unchecked(row) }

    }