Created
March 29, 2024 08:27
-
-
Save adragomir/868933a8feb077c22a26fb4d7a6798dd to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
diff --git i/Cargo.toml w/Cargo.toml | |
index 8e89e5ef3..c8e920051 100644 | |
--- i/Cargo.toml | |
+++ w/Cargo.toml | |
@@ -61,7 +61,7 @@ arrow = { version = "51.0.0", features = ["prettyprint"] } | |
arrow-array = { version = "51.0.0", default-features = false, features = ["chrono-tz"] } | |
arrow-buffer = { version = "51.0.0", default-features = false } | |
arrow-flight = { version = "51.0.0", features = ["flight-sql-experimental"] } | |
-arrow-ipc = { version = "51.0.0", default-features = false, features = ["lz4"] } | |
+arrow-ipc = { version = "51.0.0", default-features = false, features = ["lz4", "zstd"] } | |
arrow-ord = { version = "51.0.0", default-features = false } | |
arrow-schema = { version = "51.0.0", default-features = false } | |
arrow-string = { version = "51.0.0", default-features = false } | |
diff --git i/datafusion-cli/Cargo.lock w/datafusion-cli/Cargo.lock | |
index 0277d23f4..afb96b1d3 100644 | |
--- i/datafusion-cli/Cargo.lock | |
+++ w/datafusion-cli/Cargo.lock | |
@@ -257,6 +257,7 @@ dependencies = [ | |
"arrow-schema", | |
"flatbuffers", | |
"lz4_flex", | |
+ "zstd 0.13.0", | |
] | |
[[package]] | |
diff --git i/datafusion-cli/Cargo.toml w/datafusion-cli/Cargo.toml | |
index 18e143573..ec6fa925f 100644 | |
--- i/datafusion-cli/Cargo.toml | |
+++ w/datafusion-cli/Cargo.toml | |
@@ -31,7 +31,7 @@ readme = "README.md" | |
[dependencies] | |
arrow = "51.0.0" | |
-async-trait = "0.1.41" | |
+async-trait = "0.1.78" | |
aws-config = "0.55" | |
aws-credential-types = "0.55" | |
clap = { version = "3", features = ["derive", "cargo"] } | |
diff --git i/datafusion/sql/src/expr/mod.rs w/datafusion/sql/src/expr/mod.rs | |
index d1fc03194..0b97fc31d 100644 | |
--- i/datafusion/sql/src/expr/mod.rs | |
+++ w/datafusion/sql/src/expr/mod.rs | |
@@ -873,6 +873,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | |
planner_context: &mut PlannerContext, | |
) -> Result<GetFieldAccess> { | |
let field = match expr.clone() { | |
+ SQLExpr::Identifier(ident) => GetFieldAccess::NamedStructField { | |
+ name: ScalarValue::from(ident.value), | |
+ }, | |
SQLExpr::Value( | |
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s), | |
) => GetFieldAccess::NamedStructField { | |
diff --git i/datafusion/sql/tests/sql_integration.rs w/datafusion/sql/tests/sql_integration.rs | |
index 101c31039..dada3ebb4 100644 | |
--- i/datafusion/sql/tests/sql_integration.rs | |
+++ w/datafusion/sql/tests/sql_integration.rs | |
@@ -52,6 +52,104 @@ fn test_schema_support() { | |
); | |
} | |
+#[derive(Default)] | |
+struct AdrCtx { | |
+ options: ConfigOptions, | |
+ udfs: HashMap<String, Arc<ScalarUDF>>, | |
+ udafs: HashMap<String, Arc<AggregateUDF>>, | |
+} | |
+ | |
+impl AdrCtx { | |
+ fn options_mut(&mut self) -> &mut ConfigOptions { | |
+ &mut self.options | |
+ } | |
+ | |
+ fn with_udf(mut self, udf: ScalarUDF) -> Self { | |
+ self.udfs.insert(udf.name().to_string(), Arc::new(udf)); | |
+ self | |
+ } | |
+} | |
+ | |
+impl ContextProvider for AdrCtx { | |
+ fn get_table_source(&self, name: TableReference) -> Result<Arc<dyn TableSource>> { | |
+ let schema = match name.table() { | |
+ "tbl" => Ok(Schema::new(vec![ | |
+ Field::new( | |
+ "struct_field", | |
+ DataType::Struct( | |
+ Fields::from(vec![ | |
+ Field::new( | |
+ "subfield1", | |
+ DataType::List( | |
+ Arc::new( | |
+ Field::new( | |
+ "substruct1", | |
+ DataType::Struct( | |
+ Fields::from(vec![ | |
+ Field::new("subsubfield1", DataType::Int32, true), | |
+ Field::new("subsubfield2", DataType::Binary, true), | |
+ ]) | |
+ ), | |
+ true | |
+ ) | |
+ ) | |
+ ), | |
+ true | |
+ ) | |
+ ]) | |
+ ), | |
+ true | |
+ ), | |
+ ])), | |
+ _ => plan_err!("No table named: {} found", name.table()), | |
+ }; | |
+ match schema { | |
+ Ok(t) => Ok(Arc::new(EmptyTable::new(Arc::new(t)))), | |
+ Err(e) => Err(e), | |
+ } | |
+ } | |
+ | |
+ fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> { self.udfs.get(name).cloned() } | |
+ fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> { self.udafs.get(name).cloned() } | |
+ fn get_variable_type(&self, _: &[String]) -> Option<DataType> { | |
+ unimplemented!() | |
+ } | |
+ fn get_window_meta(&self, _name: &str) -> Option<Arc<WindowUDF>> { | |
+ None | |
+ } | |
+ fn options(&self) -> &ConfigOptions { | |
+ &self.options | |
+ } | |
+ fn create_cte_work_table(&self, _name: &str, schema: SchemaRef) -> Result<Arc<dyn TableSource>> { Ok(Arc::new(EmptyTable::new(schema))) } | |
+ fn udfs_names(&self) -> Vec<String> { | |
+ self.udfs.keys().cloned().collect() | |
+ } | |
+ fn udafs_names(&self) -> Vec<String> { | |
+ self.udafs.keys().cloned().collect() | |
+ } | |
+ fn udwfs_names(&self) -> Vec<String> { | |
+ Vec::new() | |
+ } | |
+} | |
+ | |
+#[test] | |
+fn test_adr_1() { | |
+ let dialect = &GenericDialect {}; | |
+ let sql = r#"SELECT struct_field["subfield1"][0]["subsubfield1"] FROM tbl;"#; | |
+ println!("sql: {}", &sql); | |
+ let result = DFParser::parse_sql_with_dialect(sql, dialect); | |
+ println!("result: {:?}", result); | |
+ let mut ast = result.unwrap(); | |
+ println!("ast: {:?}", ast); | |
+ let context = AdrCtx::default(); | |
+ let planner = SqlToRel::new(&context); | |
+ let mut parse_result = DFParser::parse_sql_with_dialect(sql, dialect).unwrap(); | |
+ let plan_result = planner.statement_to_plan(parse_result.pop_front().unwrap()); | |
+ println!("result: {:?}", plan_result); | |
+ | |
+ // planner.statement_to_plan(ast.pop_front().unwrap()) | |
+} | |
+ | |
#[test] | |
fn parse_decimals() { | |
let test_data = [ |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
patch -p1 < datafusion-struct-access.patch
in the arrow-datafusion dir