Merge branch 'main' into bigquery-approx-median

Author: ssabdb

ibis/backends/impala/__init__.py

@@ -31,7 +31,7 @@
 from ibis.backends.sql import SQLBackend
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
+    from collections.abc import Iterable, Mapping
     from pathlib import Path
     from urllib.parse import ParseResult
 
@@ -708,14 +708,14 @@ def _wrap_new_table(self, name, database):
 
     def insert(
         self,
-        name,
+        name: str,
         /,
-        obj=None,
+        obj: ir.Table | Any,
         *,
-        database=None,
-        overwrite=False,
-        partition=None,
-        validate=True,
+        database: str | tuple[str, str] | None = None,
+        overwrite: bool = False,
+        partition: Iterable[str] | dict[str, Any] | None = None,
+        validate: bool = True,
     ) -> None:
         """Insert into an Impala table.
 
@@ -733,8 +733,8 @@ def insert(
             For partitioned tables, indicate the partition that's being
             inserted into, either with an ordered list of partition keys or a
             dict of partition field name to value. For example for the
-            partition (year=2007, month=7), this can be either (2007, 7) or
-            {'year': 2007, 'month': 7}.
+            partition (year=2007, region='CA'), this can be either (2007, 'CA') or
+            {'year': 2007, 'region': 'CA'}.
         validate
             If True, do more rigorous validation that schema of table being
             inserted is compatible with the existing table
@@ -794,6 +794,7 @@ def insert(
         statement = ddl.InsertSelect(
             self._fully_qualified_name(name, database),
             self.compile(obj),
+            obj.columns,
             partition=partition,
             partition_schema=partition_schema,
             overwrite=overwrite,

ibis/backends/impala/ddl.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
 import json
+from typing import TYPE_CHECKING
 
 import sqlglot as sg
 
 import ibis.expr.schema as sch
 from ibis.backends.sql.datatypes import ImpalaType
 from ibis.backends.sql.ddl import DDL, DML, CreateDDL, DropFunction, DropObject
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 class ImpalaBase:
     dialect = "hive"
@@ -51,8 +55,7 @@ def __init__(self, name, path=None, can_exist=False):
     def compile(self):
         name = self.quote(self.name)
 
-        create_decl = "CREATE DATABASE"
-        create_line = f"{create_decl} {self._if_exists()}{name}"
+        create_line = f"CREATE DATABASE {self._if_exists()}{name}"
         if self.path is not None:
             create_line += f"\nLOCATION '{self.path}'"
 
@@ -405,15 +408,16 @@ class PartitionProperties(ImpalaBase, DDL):
 
     def __init__(
         self,
-        table,
+        # Must already be fully qualified, eg `my_db`.`my_table`
+        table_location: str,
         partition,
         partition_schema,
         location=None,
         format=None,
         tbl_properties=None,
         serde_properties=None,
     ):
-        self.table = table
+        self.table_location = table_location
         self.location = location
         self.format = self.sanitize_format(format)
         self.tbl_properties = tbl_properties
@@ -427,7 +431,7 @@ def compile(self):
             part = f"{self._command} {part}"
 
         props = self._format_properties()
-        return f"ALTER TABLE {self.table} {part}{props}"
+        return f"ALTER TABLE {self.table_location} {part}{props}"
 
     def _format_properties(self):
         tokens = []
@@ -551,6 +555,7 @@ def __init__(
         self,
         table_name,
         select_expr,
+        columns: Iterable[str],
         database=None,
         partition=None,
         partition_schema=None,
@@ -559,6 +564,7 @@ def __init__(
         self.table_name = table_name
         self.database = database
         self.select = select_expr
+        self.columns = columns
 
         self.partition = partition
         self.partition_schema = partition_schema
@@ -577,6 +583,7 @@ def compile(self):
         else:
             partition = ""
 
+        columns_str = "(" + ", ".join(self.quote(col) for col in self.columns) + ")"
         select_query = self.select
         scoped_name = self.scoped_name(self.table_name, self.database)
-        return f"{cmd} {scoped_name}{partition}\n{select_query}"
+        return f"{cmd} {scoped_name} {columns_str} {partition}\n{select_query}"

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_add_partition/out.sql

@@ -1 +1 @@
-ALTER TABLE tbl ADD PARTITION (year=2007, month=4)
+ALTER TABLE tbl ADD PARTITION (`year`=2007, `region`='CA')

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_add_partition_string_key/out.sql

@@ -1 +1 @@
-ALTER TABLE tbl ADD PARTITION (foo=5, bar="qux")
+ALTER TABLE tbl ADD PARTITION (`foo`=5, `bar`='qux')

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_add_partition_with_props/out.sql

@@ -1,2 +1,2 @@
-ALTER TABLE tbl ADD PARTITION (year=2007, month=4)
+ALTER TABLE tbl ADD PARTITION (`year`=2007, `region`='CA')
 LOCATION '/users/foo/my-data'

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_alter_partition_properties/out1.sql

@@ -1,2 +1,2 @@
-ALTER TABLE tbl PARTITION (year=2007, month=4)
+ALTER TABLE tbl PARTITION (`year`=2007, `region`='CA')
 SET LOCATION '/users/foo/my-data'

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_alter_partition_properties/out2.sql

@@ -1,2 +1,2 @@
-ALTER TABLE tbl PARTITION (year=2007, month=4)
+ALTER TABLE tbl PARTITION (`year`=2007, `region`='CA')
 SET FILEFORMAT AVRO

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_alter_partition_properties/out3.sql

@@ -1,4 +1,4 @@
-ALTER TABLE tbl PARTITION (year=2007, month=4)
+ALTER TABLE tbl PARTITION (`year`=2007, `region`='CA')
 SET TBLPROPERTIES (
   'bar'='2',
   'foo'='1'

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_alter_partition_properties/out4.sql

@@ -1,4 +1,4 @@
-ALTER TABLE tbl PARTITION (year=2007, month=4)
+ALTER TABLE tbl PARTITION (`year`=2007, `region`='CA')
 SET SERDEPROPERTIES (
   'baz'='3'
 )

ibis/backends/impala/tests/snapshots/test_ddl_compilation/test_alter_table_properties/out1.sql

@@ -1,2 +1,2 @@
-ALTER TABLE tbl PARTITION (year=2007, month=4)
+ALTER TABLE tbl PARTITION (`year`=2007, `region`='CA')
 SET LOCATION '/users/foo/my-data'