Skip to content

Instantly share code, notes, and snippets.

@adragomir
Created March 29, 2024 08:27
Show Gist options
  • Save adragomir/868933a8feb077c22a26fb4d7a6798dd to your computer and use it in GitHub Desktop.
Save adragomir/868933a8feb077c22a26fb4d7a6798dd to your computer and use it in GitHub Desktop.
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 = [
@adragomir
Copy link
Author

patch -p1 < datafusion-struct-access.patch in the arrow-datafusion dir

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment