Merge pull request #101 from embulk/support-field-partitioning

Support field partitioning

Author: sonots

README.md

@@ -158,6 +158,8 @@ This is not transactional, i.e., if fails, the target table could have some rows
 
 ```is_skip_job_result_check``` must be false when replace mode
 
+NOTE: BigQuery does not support replacing (actually, copying into) a non-partitioned table with a paritioned table atomically. You must once delete the non-partitioned table, otherwise, you get `Incompatible table partitioning specification when copying to the column partitioned table` error.
+
 ##### replace_backup
 
 1. Load to temporary table (Create and WRITE_APPEND in parallel)

example/config_delete_in_advance_field_partitioned_table.yml

@@ -0,0 +1,33 @@
+in:
+  type: file
+  path_prefix: example/example.csv
+  parser:
+    type: csv
+    charset: UTF-8
+    newline: CRLF
+    null_string: 'NULL'
+    skip_header_lines: 1
+    comment_line_marker: '#'
+    columns:
+      - {name: date,        type: string}
+      - {name: timestamp,   type: timestamp, format: "%Y-%m-%d %H:%M:%S.%N", timezone: "+09:00"}
+      - {name: "null",      type: string}
+      - {name: long,        type: long}
+      - {name: string,      type: string}
+      - {name: double,      type: double}
+      - {name: boolean,     type: boolean}
+out:
+  type: bigquery
+  mode: delete_in_advance
+  auth_method: json_key
+  json_keyfile: example/your-project-000.json
+  dataset: your_dataset_name
+  table: your_field_partitioned_table_name
+  source_format: NEWLINE_DELIMITED_JSON
+  compression: NONE
+  auto_create_dataset: true
+  auto_create_table: true
+  schema_file: example/schema.json
+  time_partitioning:
+    type: 'DAY'
+    field: timestamp

example/config_replace_backup_field_partitioned_table.yml

@@ -0,0 +1,34 @@
+in:
+  type: file
+  path_prefix: example/example.csv
+  parser:
+    type: csv
+    charset: UTF-8
+    newline: CRLF
+    null_string: 'NULL'
+    skip_header_lines: 1
+    comment_line_marker: '#'
+    columns:
+      - {name: date,        type: string}
+      - {name: timestamp,   type: timestamp, format: "%Y-%m-%d %H:%M:%S.%N", timezone: "+09:00"}
+      - {name: "null",      type: string}
+      - {name: long,        type: long}
+      - {name: string,      type: string}
+      - {name: double,      type: double}
+      - {name: boolean,     type: boolean}
+out:
+  type: bigquery
+  mode: replace_backup
+  auth_method: json_key
+  json_keyfile: example/your-project-000.json
+  dataset: your_dataset_name
+  table: your_field_partitioned_table_name
+  table_old: your_field_partitioned_table_name_old
+  source_format: NEWLINE_DELIMITED_JSON
+  compression: NONE
+  auto_create_dataset: true
+  auto_create_table: true
+  schema_file: example/schema.json
+  time_partitioning:
+    type: 'DAY'
+    field: 'timestamp'

example/config_replace_backup_paritioned_table.yml renamed to example/config_replace_backup_partitioned_table.yml

@@ -0,0 +1,33 @@
+in:
+  type: file
+  path_prefix: example/example.csv
+  parser:
+    type: csv
+    charset: UTF-8
+    newline: CRLF
+    null_string: 'NULL'
+    skip_header_lines: 1
+    comment_line_marker: '#'
+    columns:
+      - {name: date,        type: string}
+      - {name: timestamp,   type: timestamp, format: "%Y-%m-%d %H:%M:%S.%N", timezone: "+09:00"}
+      - {name: "null",      type: string}
+      - {name: long,        type: long}
+      - {name: string,      type: string}
+      - {name: double,      type: double}
+      - {name: boolean,     type: boolean}
+out:
+  type: bigquery
+  mode: replace
+  auth_method: json_key
+  json_keyfile: example/your-project-000.json
+  dataset: your_dataset_name
+  table: your_field_partitioned_table_name
+  source_format: NEWLINE_DELIMITED_JSON
+  compression: NONE
+  auto_create_dataset: true
+  auto_create_table: true
+  schema_file: example/schema.json
+  time_partitioning:
+    type: 'DAY'
+    field: 'timestamp'

example/config_replace_field_partitioned_table.yml

@@ -64,7 +64,7 @@ def self.configure(config, schema, task_count)
           'default_timestamp_format'       => config.param('default_timestamp_format',       :string,  :default => ValueConverterFactory::DEFAULT_TIMESTAMP_FORMAT),
           'payload_column'                 => config.param('payload_column',                 :string,  :default => nil),
           'payload_column_index'           => config.param('payload_column_index',           :integer, :default => nil),
-          
+
           'open_timeout_sec'               => config.param('open_timeout_sec',               :integer, :default => nil),
           'timeout_sec'                    => config.param('timeout_sec',                    :integer, :default => nil), # google-api-ruby-client < v0.11.0
           'send_timeout_sec'               => config.param('send_timeout_sec',               :integer, :default => nil), # google-api-ruby-client >= v0.11.0
@@ -276,7 +276,7 @@ def self.transaction_report(task, responses)
           sum + (response ? response.statistics.load.output_rows.to_i : 0)
         end
         if task['temp_table']
-          num_output_rows = bigquery.get_table(task['temp_table']).num_rows.to_i
+          num_output_rows = bigquery.get_table_or_partition(task['temp_table']).num_rows.to_i
         else
           num_output_rows = num_response_rows
         end
@@ -306,37 +306,48 @@ def self.auto_create(task, bigquery)
 
         case task['mode']
         when 'delete_in_advance'
-          if task['time_partitioning']
-            bigquery.delete_partition(task['table'])
-          else
-            bigquery.delete_table(task['table'])
+          bigquery.delete_partition(task['table'])
+          bigquery.create_table_if_not_exists(task['table'])
+        when 'replace'
+          bigquery.create_table_if_not_exists(task['temp_table'])
+          if Helper.has_partition_decorator?(task['table'])
+            if task['auto_create_table']
+              bigquery.create_table_if_not_exists(task['table'])
+            else
+              bigquery.get_table(task['table']) # raises NotFoundError
+            end
           end
-          bigquery.create_table(task['table'])
-        when 'replace', 'replace_backup', 'append'
-          bigquery.create_table(task['temp_table'])
-          if task['time_partitioning']
+        when 'append'
+          bigquery.create_table_if_not_exists(task['temp_table'])
+          if Helper.has_partition_decorator?(task['table'])
             if task['auto_create_table']
-              bigquery.create_table(task['table'])
+              bigquery.create_table_if_not_exists(task['table'])
             else
               bigquery.get_table(task['table']) # raises NotFoundError
             end
           end
-        else # append_direct
-          if task['auto_create_table']
-            bigquery.create_table(task['table'])
-          else
-            bigquery.get_table(task['table']) # raises NotFoundError
+        when 'replace_backup'
+          bigquery.create_table_if_not_exists(task['temp_table'])
+          if Helper.has_partition_decorator?(task['table'])
+            if task['auto_create_table']
+              bigquery.create_table_if_not_exists(task['table'])
+            else
+              bigquery.get_table(task['table']) # raises NotFoundError
+            end
           end
-        end
-
-        if task['mode'] == 'replace_backup'
-          if task['time_partitioning'] and Helper.has_partition_decorator?(task['table_old'])
+          if Helper.has_partition_decorator?(task['table_old'])
             if task['auto_create_table']
-              bigquery.create_table(task['table_old'], dataset: task['dataset_old'])
+              bigquery.create_table_if_not_exists(task['table_old'], dataset: task['dataset_old'])
             else
               bigquery.get_table(task['table_old'], dataset: task['dataset_old']) # raises NotFoundError
             end
           end
+        else # append_direct
+          if task['auto_create_table']
+            bigquery.create_table_if_not_exists(task['table'])
+          else
+            bigquery.get_table(task['table']) # raises NotFoundError
+          end
         end
       end
 
@@ -403,7 +414,7 @@ def self.transaction(config, schema, task_count, &control)
 
             if task['mode'] == 'replace_backup'
               begin
-                bigquery.get_table(task['table'])
+                bigquery.get_table_or_partition(task['table'])
                 bigquery.copy(task['table'], task['table_old'], task['dataset_old'])
               rescue NotFoundError
               end
@@ -515,7 +526,7 @@ def load_rehearsal
 
         self.class.rehearsal_thread = Thread.new do
           begin
-            bigquery.create_table(task['rehearsal_table'])
+            bigquery.create_table_if_not_exists(task['rehearsal_table'])
             response = bigquery.load(rehearsal_path, task['rehearsal_table'])
             num_output_rows = response ? response.statistics.load.output_rows.to_i : 0
             Embulk.logger.info { "embulk-output-bigquery: Loaded rehearsal #{num_output_rows}" }

example/config_replace_paritioned_table.yml renamed to example/config_replace_partitioned_table.yml

@@ -116,11 +116,11 @@ def load_from_gcs(object_uris, table)
               if @location
                 body[:job_reference][:location] = @location
               end
-              
+
               if @task['schema_update_options']
                 body[:configuration][:load][:schema_update_options] = @task['schema_update_options']
               end
-              
+
               opts = {}
 
               Embulk.logger.debug { "embulk-output-bigquery: insert_job(#{@project}, #{body}, #{opts})" }
@@ -412,7 +412,7 @@ def get_dataset(dataset = nil)
           end
         end
 
-        def create_table(table, dataset: nil, options: nil)
+        def create_table_if_not_exists(table, dataset: nil, options: nil)
           begin
             dataset ||= @dataset
             options ||= {}
@@ -466,8 +466,17 @@ def create_table(table, dataset: nil, options: nil)
         end
 
         def delete_table(table, dataset: nil)
+          table = Helper.chomp_partition_decorator(table)
+          delete_table_or_partition(table, dataset: dataset)
+        end
+
+        def delete_partition(table, dataset: nil)
+          delete_table_or_partition(table, dataset: dataset)
+        end
+
+        # if `table` with a partition decorator is given, a partition is deleted.
+        def delete_table_or_partition(table, dataset: nil)
           begin
-            table = Helper.chomp_partition_decorator(table)
             dataset ||= @dataset
             Embulk.logger.info { "embulk-output-bigquery: Delete table... #{@project}:#{dataset}.#{table}" }
             with_network_retry { client.delete_table(@project, dataset, table) }
@@ -486,8 +495,16 @@ def delete_table(table, dataset: nil)
         end
 
         def get_table(table, dataset: nil)
+          table = Helper.chomp_partition_decorator(table)
+          get_table_or_partition(table)
+        end
+
+        def get_partition(table, dataset: nil)
+          get_table_or_partition(table)
+        end
+
+        def get_table_or_partition(table, dataset: nil)
           begin
-            table = Helper.chomp_partition_decorator(table)
             dataset ||= @dataset
             Embulk.logger.info { "embulk-output-bigquery: Get table... #{@project}:#{dataset}.#{table}" }
             with_network_retry { client.get_table(@project, dataset, table) }
@@ -503,21 +520,6 @@ def get_table(table, dataset: nil)
             raise Error, "failed to get table #{@project}:#{dataset}.#{table}, response:#{response}"
           end
         end
-
-        # Is this only a way to drop partition?
-        def delete_partition(table_with_partition, dataset: nil)
-          dataset ||= @dataset
-          begin
-            table = Helper.chomp_partition_decorator(table_with_partition)
-            get_table(table, dataset: dataset)
-          rescue NotFoundError
-          else
-            Embulk.logger.info { "embulk-output-bigquery: Delete partition... #{@project}:#{dataset}.#{table_with_partition}" }
-            Tempfile.create('embulk_output_bigquery_empty_file_') do |fp|
-              load(fp.path, table_with_partition, write_disposition: 'WRITE_TRUNCATE')
-            end
-          end
-        end
       end
     end
   end

example/example.jsonl

@@ -7,12 +7,16 @@ class Bigquery < OutputPlugin
       class Helper
         PARTITION_DECORATOR_REGEXP = /\$.+\z/
 
-        def self.has_partition_decorator?(table)
-          !!(table =~ PARTITION_DECORATOR_REGEXP)
+        def self.field_partitioning?(task)
+          (task['time_partitioning'] || {}).key?('field')
         end
 
-        def self.chomp_partition_decorator(table)
-          table.sub(PARTITION_DECORATOR_REGEXP, '')
+        def self.has_partition_decorator?(table_name)
+          !!(table_name =~ PARTITION_DECORATOR_REGEXP)
+        end
+
+        def self.chomp_partition_decorator(table_name)
+          table_name.sub(PARTITION_DECORATOR_REGEXP, '')
         end
 
         def self.bq_type_from_embulk_type(embulk_type)