Skip to content

Commit

Permalink
Add Model Export support
Browse files Browse the repository at this point in the history
  • Loading branch information
quartzmo committed Aug 11, 2020
1 parent 2302328 commit 4fc030b
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 0 deletions.
39 changes: 39 additions & 0 deletions google-cloud-bigquery/acceptance/bigquery/bigquery_test.rb
Expand Up @@ -54,6 +54,22 @@
t
end
let(:dataset_with_access_id) { "#{prefix}_dataset_with_access" }
let(:model_id) { "model_#{SecureRandom.hex(4)}" }
let :model_sql do
model_sql = <<~MODEL_SQL
CREATE MODEL #{dataset.dataset_id}.#{model_id}
OPTIONS (
model_type='linear_reg',
max_iteration=1,
learn_rate=0.4,
learn_rate_strategy='constant'
) AS (
SELECT 'a' AS f1, 2.0 AS label
UNION ALL
SELECT 'b' AS f1, 3.8 AS label
)
MODEL_SQL
end

before do
dataset_2
Expand Down Expand Up @@ -230,6 +246,29 @@
_(downloaded_file.size).must_be :>, 0
end
end
focus
it "extracts a model to a GCS url with extract" do
job = dataset.query_job model_sql
job.wait_until_done!
_(job).wont_be :failed?

# can get the model
model = dataset.model model_id
_(model).must_be_kind_of Google::Cloud::Bigquery::Model

Tempfile.open "temp_extract_model" do |tmp|
dest_file_name = random_file_destination_name_model
extract_url = "gs://#{bucket.name}/#{dest_file_name}"
result = bigquery.extract model, extract_url do |j|
j.location = "US"
end
_(result).must_equal true

extract_file = bucket.file dest_file_name
downloaded_file = extract_file.download tmp.path
_(downloaded_file.size).must_be :>, 0
end
end

it "copies a readonly table to another table with copy" do
result = bigquery.copy samples_public_table, "#{dataset_id}.shakespeare_copy", create: :needed, write: :empty do |j|
Expand Down
4 changes: 4 additions & 0 deletions google-cloud-bigquery/acceptance/bigquery_helper.rb
Expand Up @@ -125,6 +125,10 @@ def random_file_destination_name
"kitten-test-data-#{SecureRandom.hex}.json"
end

def random_file_destination_name_model
"my-test-extract-model-#{SecureRandom.hex}"
end

def assert_data data
assert_equal Google::Cloud::Bigquery::Data, data.class
refute_nil data.kind
Expand Down

0 comments on commit 4fc030b

Please sign in to comment.