refactor: use reexport in arrow instead arrow_schema

Chen-Yuan-Lai · Feb 9, 2025 · 5662ab0 · 5662ab0
1 parent 3550758
commit 5662ab0
Show file tree

Hide file tree

Showing 89 changed files with 146 additions and 155 deletions.
diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs
@@ -30,10 +30,9 @@ use datafusion_expr::{BinaryExpr, Operator};
 use arrow::{
     array::{Array, ArrayRef, AsArray, StringBuilder},
     compute::{and, cast, prep_null_mask_filter},
-    datatypes::{DataType, Field, Schema},
+    datatypes::{DataType, Field, Fields, Schema},
     record_batch::RecordBatch,
 };
-use arrow_schema::Fields;
 use datafusion_expr::execution_props::ExecutionProps;
 use futures::stream::FuturesUnordered;
 use futures::{stream::BoxStream, StreamExt, TryStreamExt};

diff --git a/datafusion/common/src/column.rs b/datafusion/common/src/column.rs
@@ -20,7 +20,7 @@
 use crate::error::_schema_err;
 use crate::utils::{parse_identifiers_normalized, quote_identifier};
 use crate::{DFSchema, Diagnostic, Result, SchemaError, Spans, TableReference};
-use arrow_schema::{Field, FieldRef};
+use arrow::datatypes::{Field, FieldRef};
 use std::collections::HashSet;
 use std::convert::Infallible;
 use std::fmt;
@@ -380,8 +380,7 @@ impl fmt::Display for Column {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::datatypes::DataType;
-    use arrow_schema::SchemaBuilder;
+    use arrow::datatypes::{DataType, SchemaBuilder};
     use std::sync::Arc;
 
     fn create_qualified_schema(qualifier: &str, names: Vec<&str>) -> Result<DFSchema> {

diff --git a/datafusion/common/src/dfschema.rs b/datafusion/common/src/dfschema.rs
@@ -30,8 +30,9 @@ use crate::{
 };
 
 use arrow::compute::can_cast_types;
-use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema, SchemaRef};
-use arrow_schema::SchemaBuilder;
+use arrow::datatypes::{
+    DataType, Field, FieldRef, Fields, Schema, SchemaBuilder, SchemaRef,
+};
 
 /// A reference-counted reference to a [DFSchema].
 pub type DFSchemaRef = Arc<DFSchema>;
@@ -56,7 +57,7 @@ pub type DFSchemaRef = Arc<DFSchema>;
 ///
 /// ```rust
 /// use datafusion_common::{DFSchema, Column};
-/// use arrow_schema::{DataType, Field, Schema};
+/// use arrow::datatypes::{DataType, Field, Schema};
 ///
 /// let arrow_schema = Schema::new(vec![
 ///    Field::new("c1", DataType::Int32, false),
@@ -77,7 +78,7 @@ pub type DFSchemaRef = Arc<DFSchema>;
 ///
 /// ```rust
 /// use datafusion_common::{DFSchema, Column};
-/// use arrow_schema::{DataType, Field, Schema};
+/// use arrow::datatypes::{DataType, Field, Schema};
 ///
 /// let arrow_schema = Schema::new(vec![
 ///    Field::new("c1", DataType::Int32, false),
@@ -94,8 +95,7 @@ pub type DFSchemaRef = Arc<DFSchema>;
 ///
 /// ```rust
 /// use datafusion_common::DFSchema;
-/// use arrow_schema::Schema;
-/// use arrow::datatypes::Field;
+/// use arrow::datatypes::{Schema, Field};
 /// use std::collections::HashMap;
 ///
 /// let df_schema = DFSchema::from_unqualified_fields(vec![

diff --git a/datafusion/common/src/file_options/parquet_writer.rs b/datafusion/common/src/file_options/parquet_writer.rs
@@ -25,7 +25,7 @@ use crate::{
     DataFusionError, Result, _internal_datafusion_err,
 };
 
-use arrow_schema::Schema;
+use arrow::datatypes::Schema;
 // TODO: handle once deprecated
 #[allow(deprecated)]
 use parquet::{

diff --git a/datafusion/common/src/param_value.rs b/datafusion/common/src/param_value.rs
@@ -17,7 +17,7 @@
 
 use crate::error::{_plan_datafusion_err, _plan_err};
 use crate::{Result, ScalarValue};
-use arrow_schema::DataType;
+use arrow::datatypes::DataType;
 use std::collections::HashMap;
 
 /// The parameter value corresponding to the placeholder

diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs
@@ -55,10 +55,9 @@ use arrow::datatypes::{
     Int8Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit,
     IntervalYearMonthType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType,
     TimestampNanosecondType, TimestampSecondType, UInt16Type, UInt32Type, UInt64Type,
-    UInt8Type, DECIMAL128_MAX_PRECISION,
+    UInt8Type, UnionFields, UnionMode, DECIMAL128_MAX_PRECISION,
 };
 use arrow::util::display::{array_value_to_string, ArrayFormatter, FormatOptions};
-use arrow_schema::{UnionFields, UnionMode};
 
 use crate::format::DEFAULT_CAST_OPTIONS;
 use half::f16;
@@ -3964,9 +3963,9 @@ mod tests {
     use arrow::array::{types::Float64Type, NullBufferBuilder};
     use arrow::buffer::{Buffer, OffsetBuffer};
     use arrow::compute::{is_null, kernels};
+    use arrow::datatypes::Fields;
     use arrow::error::ArrowError;
     use arrow::util::pretty::pretty_format_columns;
-    use arrow_schema::Fields;
     use chrono::NaiveDate;
     use rand::Rng;
 

diff --git a/datafusion/common/src/scalar/struct_builder.rs b/datafusion/common/src/scalar/struct_builder.rs
@@ -20,8 +20,7 @@
 use crate::error::_internal_err;
 use crate::{Result, ScalarValue};
 use arrow::array::{ArrayRef, StructArray};
-use arrow::datatypes::{DataType, FieldRef, Fields};
-use arrow_schema::Field;
+use arrow::datatypes::{DataType, Field, FieldRef, Fields};
 use std::sync::Arc;
 
 /// Builder for [`ScalarValue::Struct`].

diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs
@@ -21,7 +21,7 @@ use std::fmt::{self, Debug, Display};
 
 use crate::{Result, ScalarValue};
 
-use arrow_schema::{DataType, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Schema, SchemaRef};
 
 /// Represents a value with a degree of certainty. `Precision` is used to
 /// propagate information the precision of statistical values.

diff --git a/datafusion/common/src/test_util.rs b/datafusion/common/src/test_util.rs
@@ -338,9 +338,9 @@ macro_rules! create_array {
 macro_rules! record_batch {
     ($(($name: expr, $type: ident, $values: expr)),*) => {
         {
-            let schema = std::sync::Arc::new(arrow_schema::Schema::new(vec![
+            let schema = std::sync::Arc::new(arrow::datatypes::Schema::new(vec![
                 $(
-                    arrow_schema::Field::new($name, arrow_schema::DataType::$type, true),
+                    arrow::datatypes::Field::new($name, arrow::datatypes::DataType::$type, true),
                 )*
             ]));
 

diff --git a/datafusion/common/src/types/field.rs b/datafusion/common/src/types/field.rs
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow_schema::{Field, Fields, UnionFields};
+use arrow::datatypes::{Field, Fields, UnionFields};
 use std::hash::{Hash, Hasher};
 use std::{ops::Deref, sync::Arc};
 

diff --git a/datafusion/common/src/types/logical.rs b/datafusion/common/src/types/logical.rs
@@ -17,7 +17,7 @@
 
 use super::NativeType;
 use crate::error::Result;
-use arrow_schema::DataType;
+use arrow::datatypes::DataType;
 use core::fmt;
 use std::{cmp::Ordering, hash::Hash, sync::Arc};
 

diff --git a/datafusion/common/src/types/native.rs b/datafusion/common/src/types/native.rs
@@ -21,7 +21,7 @@ use super::{
 };
 use crate::error::{Result, _internal_err};
 use arrow::compute::can_cast_types;
-use arrow_schema::{
+use arrow::datatypes::{
     DataType, Field, FieldRef, Fields, IntervalUnit, TimeUnit, UnionFields,
 };
 use std::{fmt::Display, sync::Arc};
@@ -126,7 +126,7 @@ pub enum NativeType {
     /// nevertheless correct).
     ///
     /// ```
-    /// # use arrow_schema::{DataType, TimeUnit};
+    /// # use arrow::datatypes::{DataType, TimeUnit};
     /// DataType::Timestamp(TimeUnit::Second, None);
     /// DataType::Timestamp(TimeUnit::Second, Some("literal".into()));
     /// DataType::Timestamp(TimeUnit::Second, Some("string".to_string().into()));

diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs
@@ -30,8 +30,7 @@ use arrow::array::{
 };
 use arrow::buffer::OffsetBuffer;
 use arrow::compute::{partition, SortColumn, SortOptions};
-use arrow::datatypes::{Field, SchemaRef};
-use arrow_schema::DataType;
+use arrow::datatypes::{DataType, Field, SchemaRef};
 use sqlparser::ast::Ident;
 use sqlparser::dialect::GenericDialect;
 use sqlparser::parser::Parser;

diff --git a/datafusion/core/benches/sql_query_with_io.rs b/datafusion/core/benches/sql_query_with_io.rs
@@ -18,7 +18,7 @@
 use std::{fmt::Write, sync::Arc, time::Duration};
 
 use arrow::array::{Int64Builder, RecordBatch, UInt64Builder};
-use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use bytes::Bytes;
 use criterion::{criterion_group, criterion_main, Criterion, SamplingMode};
 use datafusion::{

diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs
@@ -45,8 +45,7 @@ use std::sync::Arc;
 
 use arrow::array::{Array, ArrayRef, Int64Array, StringArray};
 use arrow::compute::{cast, concat};
-use arrow::datatypes::{DataType, Field};
-use arrow_schema::{Schema, SchemaRef};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::config::{CsvOptions, JsonOptions};
 use datafusion_common::{
     exec_err, not_impl_err, plan_err, Column, DFSchema, DataFusionError, ParamValues,

diff --git a/datafusion/core/src/datasource/data_source.rs b/datafusion/core/src/datasource/data_source.rs
@@ -24,7 +24,7 @@ use std::sync::Arc;
 
 use crate::datasource::physical_plan::{FileOpener, FileScanConfig};
 
-use arrow_schema::SchemaRef;
+use arrow::datatypes::SchemaRef;
 use datafusion_common::Statistics;
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::DisplayFormatType;

diff --git a/datafusion/core/src/datasource/file_format/arrow.rs b/datafusion/core/src/datasource/file_format/arrow.rs
@@ -37,11 +37,12 @@ use crate::datasource::physical_plan::{
 use crate::error::Result;
 use crate::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
 
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::error::ArrowError;
 use arrow::ipc::convert::fb_to_schema;
 use arrow::ipc::reader::FileReader;
 use arrow::ipc::writer::IpcWriteOptions;
 use arrow::ipc::{root_as_message, CompressionType};
-use arrow_schema::{ArrowError, Schema, SchemaRef};
 use datafusion_catalog::Session;
 use datafusion_common::parsers::CompressionTypeVariant;
 use datafusion_common::{

diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs
@@ -43,7 +43,7 @@ use crate::physical_plan::{
 use arrow::array::RecordBatch;
 use arrow::csv::WriterBuilder;
 use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef};
-use arrow_schema::ArrowError;
+use arrow::error::ArrowError;
 use datafusion_catalog::Session;
 use datafusion_common::config::{ConfigField, ConfigFileType, CsvOptions};
 use datafusion_common::file_options::csv_writer::CsvWriterOptions;

diff --git a/datafusion/core/src/datasource/file_format/json.rs b/datafusion/core/src/datasource/file_format/json.rs
@@ -43,11 +43,10 @@ use crate::physical_plan::{
 };
 
 use arrow::array::RecordBatch;
-use arrow::datatypes::Schema;
-use arrow::datatypes::SchemaRef;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::error::ArrowError;
 use arrow::json;
 use arrow::json::reader::{infer_json_schema_from_iterator, ValueIter};
-use arrow_schema::ArrowError;
 use datafusion_catalog::Session;
 use datafusion_common::config::{ConfigField, ConfigFileType, JsonOptions};
 use datafusion_common::file_options::json_writer::JsonWriterOptions;
@@ -438,9 +437,9 @@ mod tests {
     use crate::test::object_store::local_unpartitioned_file;
 
     use arrow::compute::concat_batches;
+    use arrow::datatypes::{DataType, Field};
     use arrow::json::ReaderBuilder;
     use arrow::util::pretty;
-    use arrow_schema::{DataType, Field};
     use datafusion_common::cast::as_int64_array;
     use datafusion_common::stats::Precision;
     use datafusion_common::{assert_batches_eq, internal_err};

diff --git a/datafusion/core/src/datasource/file_format/mod.rs b/datafusion/core/src/datasource/file_format/mod.rs
@@ -38,12 +38,12 @@ use std::sync::Arc;
 use std::task::Poll;
 
 use crate::arrow::array::RecordBatch;
-use crate::arrow::datatypes::SchemaRef;
+use crate::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef};
+use crate::arrow::error::ArrowError;
 use crate::datasource::physical_plan::{FileScanConfig, FileSinkConfig};
 use crate::error::Result;
 use crate::physical_plan::{ExecutionPlan, Statistics};
 
-use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema};
 use datafusion_catalog::Session;
 use datafusion_common::file_options::file_type::FileType;
 use datafusion_common::{internal_err, not_impl_err, GetExt};

diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs
@@ -1313,7 +1313,7 @@ mod tests {
         types::Int32Type, Array, ArrayRef, DictionaryArray, Int32Array, Int64Array,
         StringArray,
     };
-    use arrow_schema::{DataType, Field};
+    use arrow::datatypes::{DataType, Field};
     use async_trait::async_trait;
     use datafusion_common::cast::{
         as_binary_array, as_binary_view_array, as_boolean_array, as_float32_array,

diff --git a/datafusion/core/src/datasource/file_format/write/demux.rs b/datafusion/core/src/datasource/file_format/write/demux.rs
@@ -31,7 +31,7 @@ use arrow::array::{
     builder::UInt64Builder, cast::AsArray, downcast_dictionary_array, RecordBatch,
     StringArray, StructArray,
 };
-use arrow_schema::{DataType, Schema};
+use arrow::datatypes::{DataType, Schema};
 use datafusion_common::cast::{
     as_boolean_array, as_date32_array, as_date64_array, as_int32_array, as_int64_array,
     as_string_array, as_string_view_array,

diff --git a/datafusion/core/src/datasource/file_format/write/mod.rs b/datafusion/core/src/datasource/file_format/write/mod.rs
@@ -26,7 +26,7 @@ use crate::datasource::physical_plan::FileSinkConfig;
 use crate::error::Result;
 
 use arrow::array::RecordBatch;
-use arrow_schema::Schema;
+use arrow::datatypes::Schema;
 use bytes::Bytes;
 use object_store::buffered::BufWriter;
 use object_store::path::Path;

diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs
@@ -39,8 +39,7 @@ use datafusion_expr::{utils::conjunction, Expr, TableProviderFilterPushDown};
 use datafusion_expr::{SortExpr, TableType};
 use datafusion_physical_plan::{empty::EmptyExec, ExecutionPlan, Statistics};
 
-use arrow::datatypes::{DataType, Field, SchemaBuilder, SchemaRef};
-use arrow_schema::Schema;
+use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, SchemaRef};
 use datafusion_common::{
     config_datafusion_err, internal_err, plan_err, project_schema, Constraints,
     SchemaExt, ToDFSchema,
@@ -1205,8 +1204,8 @@ mod tests {
     };
     use datafusion_physical_plan::collect;
 
+    use arrow::compute::SortOptions;
     use arrow::record_batch::RecordBatch;
-    use arrow_schema::SortOptions;
     use datafusion_common::stats::Precision;
     use datafusion_common::{assert_contains, ScalarValue};
     use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator};

diff --git a/datafusion/core/src/datasource/mod.rs b/datafusion/core/src/datasource/mod.rs
@@ -47,7 +47,8 @@ pub use crate::logical_expr::TableType;
 pub use datafusion_execution::object_store;
 pub use statistics::get_statistics_with_limit;
 
-use arrow_schema::{Schema, SortOptions};
+use arrow::compute::SortOptions;
+use arrow::datatypes::Schema;
 use datafusion_common::{plan_err, Result};
 use datafusion_expr::{Expr, SortExpr};
 use datafusion_physical_expr::{expressions, LexOrdering, PhysicalSortExpr};

diff --git a/datafusion/core/src/datasource/physical_plan/arrow_file.rs b/datafusion/core/src/datasource/physical_plan/arrow_file.rs
@@ -28,8 +28,8 @@ use crate::datasource::physical_plan::{
 use crate::error::Result;
 
 use arrow::buffer::Buffer;
+use arrow::datatypes::SchemaRef;
 use arrow_ipc::reader::FileDecoder;
-use arrow_schema::SchemaRef;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{Constraints, Statistics};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
@@ -321,7 +321,7 @@ impl FileOpener for ArrowOpener {
                         footer_buf[..footer_len].try_into().unwrap(),
                     )
                     .map_err(|err| {
-                        arrow_schema::ArrowError::ParseError(format!(
+                        arrow::error::ArrowError::ParseError(format!(
                             "Unable to get root as footer: {err:?}"
                         ))
                     })?;

diff --git a/datafusion/core/src/datasource/physical_plan/file_scan_config.rs b/datafusion/core/src/datasource/physical_plan/file_scan_config.rs
@@ -36,8 +36,7 @@ use arrow::array::{
     ArrayData, ArrayRef, BufferBuilder, DictionaryArray, RecordBatch, RecordBatchOptions,
 };
 use arrow::buffer::Buffer;
-use arrow::datatypes::{ArrowNativeType, UInt16Type};
-use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use arrow::datatypes::{ArrowNativeType, DataType, Field, Schema, SchemaRef, UInt16Type};
 use datafusion_common::stats::Precision;
 use datafusion_common::{
     exec_err, ColumnStatistics, Constraints, DataFusionError, Statistics,
@@ -84,8 +83,8 @@ pub fn wrap_partition_value_in_dict(val: ScalarValue) -> ScalarValue {
 /// # Example
 /// ```
 /// # use std::sync::Arc;
-/// # use arrow_schema::Schema;
-/// use datafusion::datasource::listing::PartitionedFile;
+/// # use arrow::datatypes::Schema;
+/// # use datafusion::datasource::listing::PartitionedFile;
 /// # use datafusion::datasource::physical_plan::FileScanConfig;
 /// # use datafusion_execution::object_store::ObjectStoreUrl;
 /// # use datafusion::datasource::physical_plan::ArrowSource;

diff --git a/datafusion/core/src/datasource/physical_plan/file_stream.rs b/datafusion/core/src/datasource/physical_plan/file_stream.rs
@@ -526,7 +526,7 @@ mod tests {
     use crate::test::{make_partition, object_store::register_test_store};
 
     use crate::datasource::physical_plan::CsvSource;
-    use arrow_schema::Schema;
+    use arrow::datatypes::Schema;
     use datafusion_common::internal_err;
 
     /// Test `FileOpener` which will simulate errors during file opening or scanning

diff --git a/datafusion/core/src/datasource/physical_plan/mod.rs b/datafusion/core/src/datasource/physical_plan/mod.rs
@@ -586,7 +586,7 @@ mod tests {
         BinaryArray, BooleanArray, Float32Array, Int32Array, Int64Array, RecordBatch,
         StringArray, UInt64Array,
     };
-    use arrow_schema::{Field, Schema};
+    use arrow::datatypes::{Field, Schema};
 
     use crate::datasource::schema_adapter::{
         DefaultSchemaAdapterFactory, SchemaAdapterFactory,