From c80adc869f939f75dde9c599b734aca01a47215a Mon Sep 17 00:00:00 2001 From: Fan Yang Date: Wed, 8 Jan 2025 18:05:56 +0800 Subject: [PATCH] Add fallback for simple CTAS --- backend/executor.go | 9 ++++++++- catalog/database.go | 22 ++++++++++++++++++++++ main_test.go | 36 +++--------------------------------- 3 files changed, 33 insertions(+), 34 deletions(-) diff --git a/backend/executor.go b/backend/executor.go index 9953065..f7d103b 100644 --- a/backend/executor.go +++ b/backend/executor.go @@ -143,12 +143,19 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (sql.Row } return b.base.Build(ctx, root, r) // ResolvedTable is for `SELECT * FROM table` and `TABLE table` - // SubqueryAlias is for `SELECT * FROM view`` + // SubqueryAlias is for `SELECT * FROM view` case *plan.ResolvedTable, *plan.SubqueryAlias, *plan.TableAlias: return b.executeQuery(ctx, node, conn) case *plan.Distinct, *plan.OrderedDistinct: return b.executeQuery(ctx, node, conn) case *plan.TableCopier: + // We preserve the table schema in a best-effort manner. + // For simple `CREATE TABLE t AS SELECT * FROM t`, + // we fall back to the framework to create the table and copy the data. + // For more complex cases, we directly execute the CTAS statement in DuckDB. + if _, ok := node.Source.(*plan.ResolvedTable); ok { + return b.base.Build(ctx, root, r) + } return b.executeDML(ctx, node, conn) case sql.Expressioner: return b.executeExpressioner(ctx, node, conn) diff --git a/catalog/database.go b/catalog/database.go index 76d214f..63b2cd0 100644 --- a/catalog/database.go +++ b/catalog/database.go @@ -446,3 +446,25 @@ func (d *Database) GetCollation(ctx *sql.Context) sql.CollationID { func (d *Database) SetCollation(ctx *sql.Context, collation sql.CollationID) error { return nil } + +// CopyTableData implements sql.TableCopierDatabase interface. +func (d *Database) CopyTableData(ctx *sql.Context, sourceTable string, destinationTable string) (uint64, error) { + d.mu.Lock() + defer d.mu.Unlock() + + // Use INSERT INTO ... SELECT to copy data + sql := `INSERT INTO ` + FullTableName(d.catalog, d.name, destinationTable) + ` FROM ` + FullTableName(d.catalog, d.name, sourceTable) + + res, err := adapter.Exec(ctx, sql) + if err != nil { + return 0, ErrDuckDB.New(err) + } + + // Get count of affected rows + count, err := res.RowsAffected() + if err != nil { + return 0, ErrDuckDB.New(err) + } + + return uint64(count), nil +} diff --git a/main_test.go b/main_test.go index 78aae06..202f969 100644 --- a/main_test.go +++ b/main_test.go @@ -1110,45 +1110,15 @@ func TestCreateTable(t *testing.T) { "create_table_t1_(i_int_primary_key,_b1_blob,_b2_blob,_unique_index(b1(123),_b2(456)))", "create_table_t1_(i_int_primary_key,_b1_blob,_b2_blob,_index(b1(10)),_index(b2(20)),_index(b1(123),_b2(456)))", "create_table_t1_(i_int_primary_key,_b1_blob,_b2_blob,_index(b1(10)),_index(b2(20)),_index(b1(123),_b2(456)))", - // "CREATE_TABLE_t1_as_select_*_from_mytable", - // "CREATE_TABLE_t1_as_select_*_from_mytable", - // "CREATE_TABLE_t1_as_select_*_from_mytable#01", - // "CREATE_TABLE_t1_as_select_*_from_mytable", - // "CREATE_TABLE_t1_as_select_s,_i_from_mytable", - // "CREATE_TABLE_t1_as_select_s,_i_from_mytable", - // "CREATE_TABLE_t1_as_select_distinct_s,_i_from_mytable", - // "CREATE_TABLE_t1_as_select_distinct_s,_i_from_mytable", - // "CREATE_TABLE_t1_as_select_s,_i_from_mytable_order_by_s", - // "CREATE_TABLE_t1_as_select_s,_i_from_mytable_order_by_s", - // "CREATE_TABLE_t1_as_select_s,_sum(i)_from_mytable_group_by_s", - // "CREATE_TABLE_t1_as_select_s,_sum(i)_from_mytable_group_by_s", - // "CREATE_TABLE_t1_as_select_s,_sum(i)_from_mytable_group_by_s_having_sum(i)_>_2", - // "CREATE_TABLE_t1_as_select_s,_sum(i)_from_mytable_group_by_s_having_sum(i)_>_2", - // "CREATE_TABLE_t1_as_select_s,_i_from_mytable_order_by_s_limit_1", - // "CREATE_TABLE_t1_as_select_s,_i_from_mytable_order_by_s_limit_1", - // "CREATE_TABLE_t1_as_select_concat(\"new\",_s),_i_from_mytable", - // "CREATE_TABLE_t1_as_select_concat(\"new\",_s),_i_from_mytable", + // SUM(VARCHAR) is not supported by DuckDB + "CREATE_TABLE_t1_as_select_s,_sum(i)_from_mytable_group_by_s", + "CREATE_TABLE_t1_as_select_s,_sum(i)_from_mytable_group_by_s_having_sum(i)_>_2", "display_width_for_numeric_types", "SHOW_FULL_FIELDS_FROM_numericDisplayWidthTest;", "datetime_precision", "CREATE_TABLE_tt_(pk_int_primary_key,_d_datetime(6)_default_current_timestamp(6))", "Identifier_lengths", "table_charset_options", - // "show_create_table_t3", - // "show_create_table_t4", - // "create_table_with_select_preserves_default", - // "create_table_t1_select_*_from_a;", - // "create_table_t2_select_j_from_a;", - // "create_table_t3_select_j_as_i_from_a;", - // "create_table_t4_select_j_+_1_from_a;", - // "create_table_t5_select_a.j_from_a;", - // "create_table_t6_select_sqa.j_from_(select_i,_j_from_a)_sqa;", - // "show_create_table_t7;", - // "create_table_t8_select_*_from_(select_*_from_a)_a_join_(select_*_from_b)_b;", - // "show_create_table_t9;", - // "create_table_t11_select_sum(j)_over()_as_jj_from_a;", - // "create_table_t12_select_j_from_a_group_by_j;", - // "create_table_t13_select_*_from_c;", "event_contains_CREATE_TABLE_AS", "CREATE_EVENT_foo_ON_SCHEDULE_EVERY_1_YEAR_DO_CREATE_TABLE_bar_AS_SELECT_1;", "trigger_contains_CREATE_TABLE_AS",