You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "yjshen (via GitHub)" <gi...@apache.org> on 2023/04/03 08:27:03 UTC

[GitHub] [arrow-datafusion] yjshen commented on a diff in pull request #5835: Add primary key information to CreateMemoryTable LogicalPlan node

yjshen commented on code in PR #5835:
URL: https://github.com/apache/arrow-datafusion/pull/5835#discussion_r1155614711


##########
datafusion/sql/src/statement.rs:
##########
@@ -128,69 +128,91 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 if_not_exists,
                 or_replace,
                 ..
-            } if constraints.is_empty()
-                && table_properties.is_empty()
-                && with_options.is_empty() =>
-            {
-                match query {
-                    Some(query) => {
-                        let plan = self.query_to_plan(*query, planner_context)?;
-                        let input_schema = plan.schema();
-
-                        let plan = if !columns.is_empty() {
-                            let schema = self.build_schema(columns)?.to_dfschema_ref()?;
-                            if schema.fields().len() != input_schema.fields().len() {
-                                return Err(DataFusionError::Plan(format!(
+            } if table_properties.is_empty() && with_options.is_empty() => match query {
+                Some(query) => {
+                    let plan = self.query_to_plan(*query, planner_context)?;
+                    let input_schema = plan.schema();
+
+                    let plan = if !columns.is_empty() {
+                        let schema = self.build_schema(columns)?.to_dfschema_ref()?;
+                        if schema.fields().len() != input_schema.fields().len() {
+                            return Err(DataFusionError::Plan(format!(
                             "Mismatch: {} columns specified, but result has {} columns",
                             schema.fields().len(),
                             input_schema.fields().len()
                         )));
-                            }
-                            let input_fields = input_schema.fields();
-                            let project_exprs = schema
-                                .fields()
-                                .iter()
-                                .zip(input_fields)
-                                .map(|(field, input_field)| {
-                                    cast(
-                                        col(input_field.name()),
-                                        field.data_type().clone(),
-                                    )
+                        }
+                        let input_fields = input_schema.fields();
+                        let project_exprs = schema
+                            .fields()
+                            .iter()
+                            .zip(input_fields)
+                            .map(|(field, input_field)| {
+                                cast(col(input_field.name()), field.data_type().clone())
                                     .alias(field.name())
-                                })
-                                .collect::<Vec<_>>();
-                            LogicalPlanBuilder::from(plan.clone())
-                                .project(project_exprs)?
-                                .build()?
-                        } else {
-                            plan
-                        };
-
-                        Ok(LogicalPlan::CreateMemoryTable(CreateMemoryTable {
-                            name: self.object_name_to_table_reference(name)?,
-                            input: Arc::new(plan),
-                            if_not_exists,
-                            or_replace,
-                        }))
-                    }
+                            })
+                            .collect::<Vec<_>>();
+                        LogicalPlanBuilder::from(plan.clone())
+                            .project(project_exprs)?
+                            .build()?
+                    } else {
+                        plan
+                    };
+
+                    Ok(LogicalPlan::CreateMemoryTable(CreateMemoryTable {
+                        name: self.object_name_to_table_reference(name)?,
+                        primary_key: vec![],

Review Comment:
   Should we also consider the `TableContraint` here?



##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -1490,6 +1490,8 @@ pub struct Union {
 pub struct CreateMemoryTable {
     /// The table name
     pub name: OwnedTableReference,
+    /// The ordered list of columns in the primary key, or an empty vector if none
+    pub primary_key: Vec<Column>,

Review Comment:
   I think we might not be able to get the PK index ordering from AST.



##########
datafusion/sql/tests/integration_test.rs:
##########
@@ -199,6 +199,17 @@ fn cast_to_invalid_decimal_type() {
     }
 }
 
+#[test]
+fn plan_create_table_with_pk() {
+    let sql = "create table person (id int, name string, primary key(id))";
+    let plan = r#"
+CreateMemoryTable: Bare { table: "person" }

Review Comment:
   I think it's https://github.com/apache/arrow-datafusion/blob/main/datafusion/expr/src/logical_plan/plan.rs#L1051
   and 
   https://github.com/apache/arrow-datafusion/blob/main/datafusion/common/src/table_reference.rs#L110 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org