Skip to content

Commit

Permalink
feat: re-export arrow crates for arrow-rs feature (#309)
Browse files Browse the repository at this point in the history
  • Loading branch information
mbrobbel authored Jan 9, 2025
1 parent eb02a0e commit 7391c0f
Show file tree
Hide file tree
Showing 5 changed files with 31 additions and 20 deletions.
2 changes: 1 addition & 1 deletion examples/parquet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,12 @@ use narrow::NonNullable;

#[rustversion::attr(nightly, allow(non_local_definitions))]
fn main() {
use arrow_array::RecordBatch;
use arrow_cast::pretty;
use bytes::Bytes;
use narrow::{
array::{StructArray, VariableSizeBinary},
arrow::buffer::ScalarBuffer,
arrow_array::{self, RecordBatch},
ArrayType,
};
use parquet::arrow::{arrow_reader::ParquetRecordBatchReader, ArrowWriter};
Expand Down
7 changes: 5 additions & 2 deletions examples/tpch_lineitem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,12 @@ use arrow_array::{
},
Array, RecordBatch,
};
use arrow_schema::{DataType, Field, TimeUnit};
use chrono::{DateTime, Utc};
use narrow::{array::StructArray, ArrayType};
use narrow::{
array::StructArray,
arrow_schema::{DataType, Field, TimeUnit},
ArrayType,
};
use rand::{prelude::SmallRng, Rng, SeedableRng};
use std::sync::Arc;

Expand Down
14 changes: 7 additions & 7 deletions narrow-derive/src/enum.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1047,8 +1047,8 @@ impl<'a> Enum<'a> {
let variant_idx = (0..self.variants.len()).map(|idx| idx.to_string());
let tokens = quote! {
impl #impl_generics #narrow::arrow::UnionArrayTypeFields<#variants> for #ident #ty_generics #where_clause {
fn fields() -> ::arrow_schema::Fields {
::arrow_schema::Fields::from(vec![
fn fields() -> #narrow::arrow_schema::Fields {
#narrow::arrow_schema::Fields::from(vec![
#(
<<<#self_ident as #narrow::array::union::EnumVariant<#idx>>::Data as #narrow::array::ArrayType<<#self_ident #self_ty_generics as #narrow::array::union::EnumVariant<#idx>>::Data>>::Array<
Buffer,
Expand Down Expand Up @@ -1087,7 +1087,7 @@ impl<'a> Enum<'a> {
let (_, self_ty_generics, _) = self.generics.split_for_impl();
generics.make_where_clause().predicates.extend(
self.variant_indices().map::<WherePredicate, _>(|idx| {
parse_quote!(::std::sync::Arc<dyn ::arrow_array::Array>: From<
parse_quote!(::std::sync::Arc<dyn #narrow::arrow_array::Array>: From<
<<#self_ident #self_ty_generics as #narrow::array::union::EnumVariant<#idx>>::Data as #narrow::array::ArrayType<<#self_ident #self_ty_generics as #narrow::array::union::EnumVariant<#idx>>::Data>>::Array<
Buffer,
OffsetItem,
Expand All @@ -1101,7 +1101,7 @@ impl<'a> Enum<'a> {
let (_, ty_generics, _) = generics.split_for_impl();
let idx = self.variant_indices();
let tokens = quote! {
impl #impl_generics ::std::convert::From<#ident #ty_generics> for ::std::vec::Vec<::std::sync::Arc<dyn ::arrow_array::Array>> #where_clause {
impl #impl_generics ::std::convert::From<#ident #ty_generics> for ::std::vec::Vec<::std::sync::Arc<dyn #narrow::arrow_array::Array>> #where_clause {
fn from(value: #ident #ty_generics) -> Self {
vec![
#(
Expand Down Expand Up @@ -1129,7 +1129,7 @@ impl<'a> Enum<'a> {
.visit_generics_mut(&mut generics);
generics.make_where_clause().predicates.extend(
self.variant_indices().map::<WherePredicate, _>(|idx| {
parse_quote!(::std::sync::Arc<dyn ::arrow_array::Array>: Into<
parse_quote!(::std::sync::Arc<dyn #narrow::arrow_array::Array>: Into<
<<#self_ident #self_ty_generics as #narrow::array::union::EnumVariant<#idx>>::Data as #narrow::array::ArrayType<<#self_ident #self_ty_generics as #narrow::array::union::EnumVariant<#idx>>::Data>>::Array<
Buffer,
OffsetItem,
Expand All @@ -1149,8 +1149,8 @@ impl<'a> Enum<'a> {
.into())
});
let tokens = quote! {
impl #impl_generics ::std::iter::FromIterator<::std::sync::Arc<dyn ::arrow_array::Array>> for #ident #ty_generics #where_clause {
fn from_iter<_I: ::std::iter::IntoIterator<Item = ::std::sync::Arc<dyn ::arrow_array::Array>>>(iter: _I) -> Self {
impl #impl_generics ::std::iter::FromIterator<::std::sync::Arc<dyn #narrow::arrow_array::Array>> for #ident #ty_generics #where_clause {
fn from_iter<_I: ::std::iter::IntoIterator<Item = ::std::sync::Arc<dyn #narrow::arrow_array::Array>>>(iter: _I) -> Self {
let mut iter = iter.into_iter();
const VARIANTS: usize = #len;
Self(
Expand Down
20 changes: 10 additions & 10 deletions narrow-derive/src/struct.rs
Original file line number Diff line number Diff line change
Expand Up @@ -273,9 +273,9 @@ impl Struct<'_> {
let tokens = if matches!(self.fields, Fields::Unit) {
quote!(impl #impl_generics #narrow::arrow::StructArrayTypeFields for #ident #ty_generics #where_clause {
const NAMES: &'static [&'static str] = &[#field_name];
fn fields() -> ::arrow_schema::Fields {
::arrow_schema::Fields::from([
::std::sync::Arc::new(::arrow_schema::Field::new(#field_name, ::arrow_schema::DataType::Null, true)),
fn fields() -> #narrow::arrow_schema::Fields {
#narrow::arrow_schema::Fields::from([
::std::sync::Arc::new(#narrow::arrow_schema::Field::new(#field_name, #narrow::arrow_schema::DataType::Null, true)),
])
}
})
Expand All @@ -297,8 +297,8 @@ impl Struct<'_> {
#field_name,
)*
];
fn fields() -> ::arrow_schema::Fields {
::arrow_schema::Fields::from([
fn fields() -> #narrow::arrow_schema::Fields {
#narrow::arrow_schema::Fields::from([
#fields
])
}
Expand Down Expand Up @@ -355,7 +355,7 @@ impl Struct<'_> {

let ident = self.array_struct_ident();
let tokens = quote! {
impl #impl_generics ::std::convert::From<#ident #ty_generics> for ::std::vec::Vec<::std::sync::Arc<dyn ::arrow_array::Array>> #where_clause {
impl #impl_generics ::std::convert::From<#ident #ty_generics> for ::std::vec::Vec<::std::sync::Arc<dyn #narrow::arrow_array::Array>> #where_clause {
fn from(value: #ident #ty_generics) -> Self {
vec![
#field_arrays
Expand All @@ -381,7 +381,7 @@ impl Struct<'_> {
.make_where_clause()
.predicates
.extend(self.where_predicate_fields(parse_quote!(
::std::convert::From<::std::sync::Arc<dyn ::arrow_array::Array>>
::std::convert::From<::std::sync::Arc<dyn #narrow::arrow_array::Array>>
)));
let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

Expand Down Expand Up @@ -410,8 +410,8 @@ impl Struct<'_> {
});
let ident = self.array_struct_ident();
let tokens = quote! {
impl #impl_generics ::std::convert::From<::std::vec::Vec<::std::sync::Arc<dyn ::arrow_array::Array>>> for #ident #ty_generics #where_clause {
fn from(value: ::std::vec::Vec<::std::sync::Arc<dyn ::arrow_array::Array>>) -> Self {
impl #impl_generics ::std::convert::From<::std::vec::Vec<::std::sync::Arc<dyn #narrow::arrow_array::Array>>> for #ident #ty_generics #where_clause {
fn from(value: ::std::vec::Vec<::std::sync::Arc<dyn #narrow::arrow_array::Array>>) -> Self {
let mut arrays = value.into_iter();
let result = Self #field_arrays;
assert!(arrays.next().is_none());
Expand Down Expand Up @@ -979,7 +979,7 @@ impl Struct<'_> {
parse_quote!(
<#ty as #narrow::array::ArrayType<#ty_drop>>::Array<Buffer, #narrow::offset::NA, #narrow::array::union::NA>:
::std::convert::Into<
::std::sync::Arc<dyn ::arrow_array::Array>
::std::sync::Arc<dyn #narrow::arrow_array::Array>
>
)})
}
Expand Down
8 changes: 8 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,14 @@ pub mod logical;
#[cfg(feature = "arrow-rs")]
pub mod arrow;

// Re-export arrow crates.
#[cfg(feature = "arrow-rs")]
pub use arrow_array;
#[cfg(feature = "arrow-rs")]
pub use arrow_buffer;
#[cfg(feature = "arrow-rs")]
pub use arrow_schema;

// Re-export `narrow_derive` macros when the `derive` feature is enabled.
#[cfg(feature = "derive")]
pub use narrow_derive::ArrayType;
Expand Down

0 comments on commit 7391c0f

Please sign in to comment.