Skip to content

Commit

Permalink
style: fmt
Browse files Browse the repository at this point in the history
  • Loading branch information
@timvw
timvw committed Apr 23, 2024
1 parent 5d7b6f4 commit 8bcd556
Show file tree
Hide file tree
Showing 2 changed files with 100 additions and 38 deletions.
9 changes: 5 additions & 4 deletions examples/demo.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,10 @@
// specific language governing permissions and limitations
// under the License.

use std::fmt::Debug;
use aws_config::BehaviorVersion;
use aws_sdk_glue::config::{Credentials, ProvideCredentials};
use aws_types::SdkConfig;
use dashmap::DashMap;
use datafusion::arrow::array::StringArray;
use datafusion::common::{DataFusionError, Result};
use datafusion::datasource::object_store::ObjectStoreRegistry;
Expand All @@ -28,9 +29,8 @@ use datafusion_catalogprovider_glue::catalog_provider::glue::{
};
use object_store::aws::AmazonS3Builder;
use object_store::ObjectStore;
use std::fmt::Debug;
use std::sync::Arc;
use aws_config::BehaviorVersion;
use dashmap::DashMap;
use url::Url;

#[tokio::main]
Expand All @@ -41,7 +41,8 @@ async fn main() -> Result<()> {
// Register an object store provider which creates instances for each requested s3://bucket using the sdk_config credentials
// As an alternative you can also manually register the required object_store(s)
let object_store_provider = DemoS3ObjectStoreProvider::new(&sdk_config).await?;
let runtime_config = RuntimeConfig::default().with_object_store_registry(Arc::new(object_store_provider));
let runtime_config =
RuntimeConfig::default().with_object_store_registry(Arc::new(object_store_provider));
let config = SessionConfig::new().with_information_schema(true);
let runtime = Arc::new(RuntimeEnv::new(runtime_config)?);
let ctx = SessionContext::new_with_config_rt(config, runtime);
Expand Down
129 changes: 95 additions & 34 deletions src/catalog_provider/glue.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,13 @@

use crate::error::*;
use crate::glue_data_type_parser::*;
use aws_config::BehaviorVersion;
use aws_sdk_glue::types::{Column, StorageDescriptor, Table};
use aws_sdk_glue::Client;
use aws_types::SdkConfig;
use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef, TimeUnit};
use datafusion::catalog::CatalogProvider;
use datafusion::catalog::schema::{MemorySchemaProvider, SchemaProvider};
use datafusion::catalog::CatalogProvider;
use datafusion::datasource::file_format::avro::AvroFormat;
use datafusion::datasource::file_format::csv::CsvFormat;
use datafusion::datasource::file_format::json::JsonFormat;
Expand All @@ -20,7 +21,6 @@ use datafusion::execution::context::SessionState;
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
use aws_config::BehaviorVersion;

/// Options to register a table
pub enum TableRegistrationOptions {
Expand Down Expand Up @@ -146,17 +146,11 @@ impl GlueCatalogProvider {
table_registration_options: &TableRegistrationOptions,
ctx: &SessionState,
) -> Result<Vec<Result<()>>> {
let glue_databases = self
.client
.get_databases()
.send()
.await?
.database_list;
let glue_databases = self.client.get_databases().send().await?.database_list;

let mut results = Vec::new();
for glue_database in glue_databases {
let database = glue_database
.name;
let database = glue_database.name;
let glue_tables = self
.client
.get_tables()
Expand Down Expand Up @@ -452,9 +446,7 @@ impl GlueCatalogProvider {
}

fn map_glue_column_to_arrow_field(glue_column: &Column) -> Result<Field> {
let name = glue_column
.name
.clone();
let name = glue_column.name.clone();
let glue_type = glue_column
.r#type
.as_ref()
Expand Down Expand Up @@ -513,7 +505,11 @@ mod tests {
fn test_map_tinyint_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("tinyint").build().unwrap()
&Column::builder()
.name("id")
.r#type("tinyint")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Int8, true)
Expand All @@ -525,7 +521,11 @@ mod tests {
fn test_map_smallint_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("smallint").build().unwrap()
&Column::builder()
.name("id")
.r#type("smallint")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Int16, true)
Expand All @@ -549,7 +549,11 @@ mod tests {
fn test_map_integer_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("integer").build().unwrap()
&Column::builder()
.name("id")
.r#type("integer")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Int32, true)
Expand All @@ -561,7 +565,11 @@ mod tests {
fn test_map_bigint_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("bigint").build().unwrap()
&Column::builder()
.name("id")
.r#type("bigint")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Int64, true)
Expand All @@ -573,7 +581,11 @@ mod tests {
fn test_map_float_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("float").build().unwrap()
&Column::builder()
.name("id")
.r#type("float")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Float32, true)
Expand All @@ -585,7 +597,11 @@ mod tests {
fn test_map_double_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("double").build().unwrap()
&Column::builder()
.name("id")
.r#type("double")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Float64, true)
Expand All @@ -597,7 +613,11 @@ mod tests {
fn test_map_boolean_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("boolean").build().unwrap()
&Column::builder()
.name("id")
.r#type("boolean")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Boolean, true)
Expand All @@ -609,7 +629,11 @@ mod tests {
fn test_map_binary_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("binary").build().unwrap()
&Column::builder()
.name("id")
.r#type("binary")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Binary, true)
Expand All @@ -633,7 +657,11 @@ mod tests {
fn test_map_timestamp_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("timestamp").build().unwrap()
&Column::builder()
.name("id")
.r#type("timestamp")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Timestamp(TimeUnit::Nanosecond, None), true)
Expand All @@ -645,7 +673,11 @@ mod tests {
fn test_map_string_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("string").build().unwrap()
&Column::builder()
.name("id")
.r#type("string")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Utf8, true)
Expand All @@ -669,7 +701,11 @@ mod tests {
fn test_map_varchar_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("varchar").build().unwrap()
&Column::builder()
.name("id")
.r#type("varchar")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Utf8, true)
Expand All @@ -681,7 +717,11 @@ mod tests {
fn test_map_decimal_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("decimal(12,9)").build().unwrap()
&Column::builder()
.name("id")
.r#type("decimal(12,9)")
.build()
.unwrap()
)
.unwrap(),
Field::new("id", DataType::Decimal256(12, 9), true)
Expand All @@ -693,7 +733,11 @@ mod tests {
fn test_map_array_of_bigint_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("array<bigint>").build().unwrap()
&Column::builder()
.name("id")
.r#type("array<bigint>")
.build()
.unwrap()
)
.unwrap(),
Field::new(
Expand All @@ -709,7 +753,11 @@ mod tests {
fn test_map_array_of_int_glue_column_to_arrow_field() -> Result<()> {
assert_eq!(
GlueCatalogProvider::map_glue_column_to_arrow_field(
&Column::builder().name("id").r#type("array<int>").build().unwrap()
&Column::builder()
.name("id")
.r#type("array<int>")
.build()
.unwrap()
)
.unwrap(),
Field::new(
Expand All @@ -728,7 +776,8 @@ mod tests {
&Column::builder()
.name("id")
.r#type("array<array<string>>")
.build().unwrap()
.build()
.unwrap()
)
.unwrap(),
Field::new(
Expand All @@ -751,7 +800,8 @@ mod tests {
&Column::builder()
.name("id")
.r#type("struct<reply_id:int,next_id:int>")
.build().unwrap()
.build()
.unwrap()
)
.unwrap(),
Field::new(
Expand All @@ -776,14 +826,19 @@ mod tests {
&Column::builder()
.name("id")
.r#type("struct<reply:struct<reply_id:int>>")
.build().unwrap()
.build()
.unwrap()
)
.unwrap(),
Field::new(
"id",
DataType::Struct(Fields::from(vec![Field::new(
"reply",
DataType::Struct(Fields::from(vec![Field::new("reply_id", DataType::Int32, true),])),
DataType::Struct(Fields::from(vec![Field::new(
"reply_id",
DataType::Int32,
true
),])),
true
),])),
true
Expand All @@ -799,7 +854,8 @@ mod tests {
&Column::builder()
.name("id")
.r#type("map<string,boolean>")
.build().unwrap()
.build()
.unwrap()
)
.unwrap(),
Field::new(
Expand Down Expand Up @@ -829,7 +885,8 @@ mod tests {
&Column::builder()
.name("id")
.r#type("map<string,map<string,boolean>>")
.build().unwrap()
.build()
.unwrap()
)
.unwrap(),
Field::new(
Expand Down Expand Up @@ -947,7 +1004,11 @@ mod tests {
)
);

let struct_of_int = DataType::Struct(Fields::from(vec![Field::new("reply_id", DataType::Int32, true)]));
let struct_of_int = DataType::Struct(Fields::from(vec![Field::new(
"reply_id",
DataType::Int32,
true,
)]));

// struct type
assert_eq!(
Expand Down

0 comments on commit 8bcd556

Please sign in to comment.