Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(type): improve the Fields derive macro #14934

Merged
merged 4 commits into from
Feb 5, 2024
Merged

feat(type): improve the Fields derive macro #14934

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4a0b3bc
add primary key and into_owned_row to `Fields` trait
wangrunji0408 Feb 1, 2024
c2c3b44
fix `r#` in field name
wangrunji0408 Feb 1, 2024
d801f82
support `ToOwnedDatum` for `Vec<Option<F32>>`
wangrunji0408 Feb 4, 2024
74cd287
Merge remote-tracking branch 'origin/main' into wrj/derive-fields
wangrunji0408 Feb 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 30 additions & 1 deletion src/common/fields-derive/src/gen/test_output.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,36 @@ impl ::risingwave_common::types::Fields for Data {
::risingwave_common::types::WithDataType > ::default_data_type()), ("v3", <
bool as ::risingwave_common::types::WithDataType > ::default_data_type()),
("v4", < Serial as ::risingwave_common::types::WithDataType >
::default_data_type())
::default_data_type()), ("type", < i32 as
::risingwave_common::types::WithDataType > ::default_data_type())
]
}
fn into_owned_row(self) -> ::risingwave_common::row::OwnedRow {
::risingwave_common::row::OwnedRow::new(
vec![
::risingwave_common::types::ToOwnedDatum::to_owned_datum(self.v1),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(self.v2),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(self.v3),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(self.v4),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(self.r#type)
],
)
}
fn primary_key() -> &'static [usize] {
&[1usize, 0usize]
}
}
impl From<Data> for ::risingwave_common::types::ScalarImpl {
fn from(v: Data) -> Self {
::risingwave_common::types::StructValue::new(
vec![
::risingwave_common::types::ToOwnedDatum::to_owned_datum(v.v1),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(v.v2),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(v.v3),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(v.v4),
::risingwave_common::types::ToOwnedDatum::to_owned_datum(v.r#type)
],
)
.into()
}
}
108 changes: 83 additions & 25 deletions src/common/fields-derive/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,9 @@

use proc_macro2::TokenStream;
use quote::quote;
use syn::{Data, DeriveInput, Field, Result};
use syn::{Data, DeriveInput, Result};

#[proc_macro_derive(Fields)]
#[proc_macro_derive(Fields, attributes(primary_key))]
pub fn fields(tokens: proc_macro::TokenStream) -> proc_macro::TokenStream {
inner(tokens.into()).into()
}
Expand All @@ -31,51 +31,107 @@ fn inner(tokens: TokenStream) -> TokenStream {
fn gen(tokens: TokenStream) -> Result<TokenStream> {
let input: DeriveInput = syn::parse2(tokens)?;

let DeriveInput {
attrs: _attrs,
vis: _vis,
ident,
generics,
data,
} = input;
if !generics.params.is_empty() {
let ident = &input.ident;
if !input.generics.params.is_empty() {
return Err(syn::Error::new_spanned(
generics,
input.generics,
"generics are not supported",
));
}

let Data::Struct(r#struct) = data else {
return Err(syn::Error::new_spanned(ident, "only structs are supported"));
let Data::Struct(struct_) = &input.data else {
return Err(syn::Error::new_spanned(
input.ident,
"only structs are supported",
));
};

let fields_rs = r#struct.fields;
let fields_rw: Vec<TokenStream> = fields_rs
.into_iter()
.map(|field_rs| {
let Field {
// We can support #[field(ignore)] or other useful attributes here.
attrs: _attrs,
ident: name,
ty,
..
} = field_rs;
let name = name.map_or("".to_string(), |name| name.to_string());
let fields_rw: Vec<TokenStream> = struct_
.fields
.iter()
.map(|field| {
let mut name = field.ident.as_ref().expect("field no name").to_string();
// strip leading `r#`
if name.starts_with("r#") {
name = name[2..].to_string();
}
let ty = &field.ty;
quote! {
(#name, <#ty as ::risingwave_common::types::WithDataType>::default_data_type())
}
})
.collect();
let names = struct_
.fields
.iter()
.map(|field| field.ident.as_ref().expect("field no name"))
.collect::<Vec<_>>();
let primary_key = get_primary_key(&input).map(|indices| {
quote! {
fn primary_key() -> &'static [usize] {
&[#(#indices),*]
}
}
});

Ok(quote! {
impl ::risingwave_common::types::Fields for #ident {
fn fields() -> Vec<(&'static str, ::risingwave_common::types::DataType)> {
vec![#(#fields_rw),*]
}
fn into_owned_row(self) -> ::risingwave_common::row::OwnedRow {
::risingwave_common::row::OwnedRow::new(vec![#(
::risingwave_common::types::ToOwnedDatum::to_owned_datum(self.#names)
),*])
}
#primary_key
}
impl From<#ident> for ::risingwave_common::types::ScalarImpl {
fn from(v: #ident) -> Self {
::risingwave_common::types::StructValue::new(vec![#(
::risingwave_common::types::ToOwnedDatum::to_owned_datum(v.#names)
),*]).into()
}
}
})
}

/// Get primary key indices from `#[primary_key]` attribute.
fn get_primary_key(input: &syn::DeriveInput) -> Option<Vec<usize>> {
let syn::Data::Struct(struct_) = &input.data else {
return None;
};
// find `#[primary_key(k1, k2, ...)]` on struct
let composite = input.attrs.iter().find_map(|attr| match &attr.meta {
syn::Meta::List(list) if list.path.is_ident("primary_key") => Some(&list.tokens),
_ => None,
});
if let Some(keys) = composite {
let index = |name: &str| {
struct_
.fields
.iter()
.position(|f| f.ident.as_ref().map_or(false, |i| i == name))
.expect("primary key not found")
};
return Some(
keys.to_string()
.split(',')
.map(|s| index(s.trim()))
.collect(),
);
}
// find `#[primary_key]` on fields
for (i, field) in struct_.fields.iter().enumerate() {
for attr in &field.attrs {
if matches!(&attr.meta, syn::Meta::Path(path) if path.is_ident("primary_key")) {
return Some(vec![i]);
}
}
}
None
}

#[cfg(test)]
mod tests {
use indoc::indoc;
Expand All @@ -91,11 +147,13 @@ mod tests {
fn test_gen() {
let code = indoc! {r#"
#[derive(Fields)]
#[primary_key(v2, v1)]
struct Data {
v1: i16,
v2: std::primitive::i32,
v3: bool,
v4: Serial,
r#type: i32,
}
"#};

Expand Down
65 changes: 61 additions & 4 deletions src/common/src/types/fields.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,70 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use super::DataType;
use crate::row::OwnedRow;
use crate::util::chunk_coalesce::DataChunkBuilder;

/// A struct can implements `Fields` when if can be represented as a relational Row.
///
/// Can be automatically derived with [`#[derive(Fields)]`](derive@super::Fields).
/// # Derivable
///
/// This trait can be automatically derived with [`#[derive(Fields)]`](derive@super::Fields).
/// Type of the fields must implement [`WithDataType`](super::WithDataType) and [`ToOwnedDatum`](super::ToOwnedDatum).
///
/// ```
/// # use risingwave_common::types::Fields;
///
/// #[derive(Fields)]
/// struct Data {
/// v1: i16,
/// v2: i32,
/// }
/// ```
///
/// You can add `#[primary_key]` attribute to one of the fields to specify the primary key of the table.
///
/// ```
/// # use risingwave_common::types::Fields;
///
/// #[derive(Fields)]
/// struct Data {
/// #[primary_key]
/// v1: i16,
/// v2: i32,
/// }
/// ```
///
/// If the primary key is composite, you can add `#[primary_key(...)]` attribute to the struct to specify the order of the fields.
///
/// ```
/// # use risingwave_common::types::Fields;
///
/// #[derive(Fields)]
/// #[primary_key(v2, v1)]
/// struct Data {
/// v1: i16,
/// v2: i32,
/// }
/// ```
pub trait Fields {
/// When the struct being converted to an [`Row`](crate::row::Row) or a [`DataChunk`](crate::array::DataChunk), it schema must be consistent with the `fields` call.
/// Return the schema of the struct.
fn fields() -> Vec<(&'static str, DataType)>;

/// Convert the struct to an `OwnedRow`.
fn into_owned_row(self) -> OwnedRow;

/// The primary key of the table.
fn primary_key() -> &'static [usize] {
&[]
}

/// Create a [`DataChunkBuilder`](crate::util::chunk_coalesce::DataChunkBuilder) with the schema of the struct.
fn data_chunk_builder(capacity: usize) -> DataChunkBuilder {
DataChunkBuilder::new(
Self::fields().into_iter().map(|(_, ty)| ty).collect(),
capacity,
)
}
}

#[cfg(test)]
Expand Down Expand Up @@ -48,8 +105,8 @@ mod tests {
v7: Vec<u8>,
v8: std::vec::Vec<i16>,
v9: Option<Vec<i64>>,
v10: std::option::Option<Vec<Option<f32>>>,
v11: Box<Timestamp>,
v10: std::option::Option<Vec<Option<F32>>>,
v11: Timestamp,
v14: Sub,
}

Expand Down
41 changes: 39 additions & 2 deletions src/common/src/types/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -570,10 +570,17 @@ pub trait ToOwnedDatum {
fn to_owned_datum(self) -> Datum;
}

impl ToOwnedDatum for DatumRef<'_> {
impl<T: Into<ScalarImpl>> ToOwnedDatum for T {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TBH I'd prefer introducing a new trait for doing such conversion. By ToOwnedDatum I suppose it's only for the conversion from the ref type to an owned type, just like the std's ToOwned trait with GAT support.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But logically types impl Into<ScalarImpl> can be converted ToOwnedDatum. I think it would be convenient to reuse this trait.

#[inline(always)]
fn to_owned_datum(self) -> Datum {
self.map(ScalarRefImpl::into_scalar_impl)
Some(self.into())
}
}

impl<T: Into<ScalarImpl>> ToOwnedDatum for Option<T> {
#[inline(always)]
fn to_owned_datum(self) -> Datum {
self.map(Into::into)
}
}

Expand Down Expand Up @@ -765,6 +772,36 @@ impl From<JsonbRef<'_>> for ScalarImpl {
}
}

impl<T: PrimitiveArrayItemType> From<Vec<T>> for ScalarImpl {
fn from(v: Vec<T>) -> Self {
Self::List(v.into_iter().collect())
}
}

impl<T: PrimitiveArrayItemType> From<Vec<Option<T>>> for ScalarImpl {
fn from(v: Vec<Option<T>>) -> Self {
Self::List(v.into_iter().collect())
}
}

impl From<Vec<String>> for ScalarImpl {
fn from(v: Vec<String>) -> Self {
Self::List(v.iter().map(|s| s.as_str()).collect())
}
}

impl From<Vec<u8>> for ScalarImpl {
fn from(v: Vec<u8>) -> Self {
Self::Bytea(v.into())
}
}

impl From<Bytes> for ScalarImpl {
fn from(v: Bytes) -> Self {
Self::Bytea(v.as_ref().into())
}
}

impl ScalarImpl {
/// Creates a scalar from binary.
pub fn from_binary(bytes: &Bytes, data_type: &DataType) -> Result<Self, BoxedError> {
Expand Down
Loading