Skip to content

Commit

Permalink
Migrate Array function Documentation to code (#12948)
Browse files Browse the repository at this point in the history 
* Migrate Array

* p fix

* Update datafusion/functions-nested/src/make_array.rs

Co-authored-by: Andrew Lamb <[email protected]>

* Update docs

---------

Co-authored-by: Andrew Lamb <[email protected]>
  • Loading branch information
@jonathanc-n @alamb
jonathanc-n and alamb authored Oct 16, 2024
1 parent caeabc1 commit 91899d4
Show file tree
Hide file tree
Showing 23 changed files with 2,996 additions and 1,499 deletions.
114 changes: 112 additions & 2 deletions datafusion/functions-nested/src/array_has.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,17 @@ use arrow_buffer::BooleanBuffer;
use datafusion_common::cast::as_generic_list_array;
use datafusion_common::utils::string_utils::string_array_to_vec;
use datafusion_common::{exec_err, Result, ScalarValue};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
use datafusion_expr::scalar_doc_sections::DOC_SECTION_ARRAY;
use datafusion_expr::{
ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
};
use datafusion_physical_expr_common::datum::compare_with_eq;
use itertools::Itertools;

use crate::utils::make_scalar_function;

use std::any::Any;
use std::sync::Arc;
use std::sync::{Arc, OnceLock};

// Create static instances of ScalarUDFs for each function
make_udf_expr_and_func!(ArrayHas,
Expand Down Expand Up @@ -129,6 +132,43 @@ impl ScalarUDFImpl for ArrayHas {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn documentation(&self) -> Option<&Documentation> {
Some(get_array_has_doc())
}
}

static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new();

fn get_array_has_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_ARRAY)
.with_description(
"Returns true if the array contains the element.",
)
.with_syntax_example("array_has(array, element)")
.with_sql_example(
r#"```sql
> select array_has([1, 2, 3], 2);
+-----------------------------+
| array_has(List([1,2,3]), 2) |
+-----------------------------+
| true |
+-----------------------------+
```"#,
)
.with_argument(
"array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.with_argument(
"element",
"Scalar or Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.build()
.unwrap()
})
}

fn array_has_inner_for_scalar(
Expand Down Expand Up @@ -289,6 +329,41 @@ impl ScalarUDFImpl for ArrayHasAll {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn documentation(&self) -> Option<&Documentation> {
Some(get_array_has_all_doc())
}
}

fn get_array_has_all_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_ARRAY)
.with_description(
"Returns true if all elements of sub-array exist in array.",
)
.with_syntax_example("array_has_all(array, sub-array)")
.with_sql_example(
r#"```sql
> select array_has_all([1, 2, 3, 4], [2, 3]);
+--------------------------------------------+
| array_has_all(List([1,2,3,4]), List([2,3])) |
+--------------------------------------------+
| true |
+--------------------------------------------+
```"#,
)
.with_argument(
"array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.with_argument(
"sub-array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.build()
.unwrap()
})
}

#[derive(Debug)]
Expand Down Expand Up @@ -335,6 +410,41 @@ impl ScalarUDFImpl for ArrayHasAny {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn documentation(&self) -> Option<&Documentation> {
Some(get_array_has_any_doc())
}
}

fn get_array_has_any_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_ARRAY)
.with_description(
"Returns true if any elements exist in both arrays.",
)
.with_syntax_example("array_has_any(array, sub-array)")
.with_sql_example(
r#"```sql
> select array_has_any([1, 2, 3], [3, 4]);
+------------------------------------------+
| array_has_any(List([1,2,3]), List([3,4])) |
+------------------------------------------+
| true |
+------------------------------------------+
```"#,
)
.with_argument(
"array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.with_argument(
"sub-array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.build()
.unwrap()
})
}

/// Represents the type of comparison for array_has.
Expand Down
40 changes: 37 additions & 3 deletions datafusion/functions-nested/src/cardinality.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,13 @@ use arrow_schema::DataType::{FixedSizeList, LargeList, List, Map, UInt64};
use datafusion_common::cast::{as_large_list_array, as_list_array, as_map_array};
use datafusion_common::Result;
use datafusion_common::{exec_err, plan_err};
use datafusion_expr::scalar_doc_sections::DOC_SECTION_ARRAY;
use datafusion_expr::{
ArrayFunctionSignature, ColumnarValue, ScalarUDFImpl, Signature, TypeSignature,
Volatility,
ArrayFunctionSignature, ColumnarValue, Documentation, ScalarUDFImpl, Signature,
TypeSignature, Volatility,
};
use std::any::Any;
use std::sync::Arc;
use std::sync::{Arc, OnceLock};

make_udf_expr_and_func!(
Cardinality,
Expand Down Expand Up @@ -89,6 +90,39 @@ impl ScalarUDFImpl for Cardinality {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn documentation(&self) -> Option<&Documentation> {
Some(get_cardinality_doc())
}
}

static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new();

fn get_cardinality_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_ARRAY)
.with_description(
"Returns the total number of elements in the array.",
)
.with_syntax_example("cardinality(array)")
.with_sql_example(
r#"```sql
> select cardinality([[1, 2, 3, 4], [5, 6, 7, 8]]);
+--------------------------------------+
| cardinality(List([1,2,3,4,5,6,7,8])) |
+--------------------------------------+
| 8 |
+--------------------------------------+
```"#,
)
.with_argument(
"array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.build()
.unwrap()
})
}

/// Cardinality SQL function
Expand Down
115 changes: 112 additions & 3 deletions datafusion/functions-nested/src/concat.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,8 @@

//! [`ScalarUDFImpl`] definitions for `array_append`, `array_prepend` and `array_concat` functions.
use std::{any::Any, cmp::Ordering, sync::Arc};
use std::sync::{Arc, OnceLock};
use std::{any::Any, cmp::Ordering};

use arrow::array::{Capacities, MutableArrayData};
use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
Expand All @@ -27,9 +28,10 @@ use datafusion_common::Result;
use datafusion_common::{
cast::as_generic_list_array, exec_err, not_impl_err, plan_err, utils::list_ndims,
};
use datafusion_expr::scalar_doc_sections::DOC_SECTION_ARRAY;
use datafusion_expr::{
type_coercion::binary::get_wider_type, ColumnarValue, ScalarUDFImpl, Signature,
Volatility,
type_coercion::binary::get_wider_type, ColumnarValue, Documentation, ScalarUDFImpl,
Signature, Volatility,
};

use crate::utils::{align_array_dimensions, check_datatypes, make_scalar_function};
Expand Down Expand Up @@ -91,6 +93,43 @@ impl ScalarUDFImpl for ArrayAppend {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn documentation(&self) -> Option<&Documentation> {
Some(get_array_append_doc())
}
}

static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new();

fn get_array_append_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_ARRAY)
.with_description(
"Appends an element to the end of an array.",
)
.with_syntax_example("array_append(array, element)")
.with_sql_example(
r#"```sql
> select array_append([1, 2, 3], 4);
+--------------------------------------+
| array_append(List([1,2,3]),Int64(4)) |
+--------------------------------------+
| [1, 2, 3, 4] |
+--------------------------------------+
```"#,
)
.with_argument(
"array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.with_argument(
"element",
"Element to append to the array.",
)
.build()
.unwrap()
})
}

make_udf_expr_and_func!(
Expand Down Expand Up @@ -150,6 +189,41 @@ impl ScalarUDFImpl for ArrayPrepend {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn documentation(&self) -> Option<&Documentation> {
Some(get_array_prepend_doc())
}
}

fn get_array_prepend_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_ARRAY)
.with_description(
"Prepends an element to the beginning of an array.",
)
.with_syntax_example("array_prepend(element, array)")
.with_sql_example(
r#"```sql
> select array_prepend(1, [2, 3, 4]);
+---------------------------------------+
| array_prepend(Int64(1),List([2,3,4])) |
+---------------------------------------+
| [1, 2, 3, 4] |
+---------------------------------------+
```"#,
)
.with_argument(
"element",
"Element to prepend to the array.",
)
.with_argument(
"array",
"Array expression. Can be a constant, column, or function, and any combination of array operators.",
)
.build()
.unwrap()
})
}

make_udf_expr_and_func!(
Expand Down Expand Up @@ -233,6 +307,41 @@ impl ScalarUDFImpl for ArrayConcat {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn documentation(&self) -> Option<&Documentation> {
Some(get_array_concat_doc())
}
}

fn get_array_concat_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_ARRAY)
.with_description(
"Concatenates arrays.",
)
.with_syntax_example("array_concat(array[, ..., array_n])")
.with_sql_example(
r#"```sql
> select array_concat([1, 2], [3, 4], [5, 6]);
+---------------------------------------------------+
| array_concat(List([1,2]),List([3,4]),List([5,6])) |
+---------------------------------------------------+
| [1, 2, 3, 4, 5, 6] |
+---------------------------------------------------+
```"#,
)
.with_argument(
"array",
"Array expression to concatenate. Can be a constant, column, or function, and any combination of array operators.",
)
.with_argument(
"array_n",
"Subsequent array column or literal array to concatenate.",
)
.build()
.unwrap()
})
}

/// Array_concat/Array_cat SQL function
Expand Down
Loading

0 comments on commit 91899d4

Please sign in to comment.