Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,5 @@
- https://storage.googleapis.com/<package_path>
transforms:
CopyFilesToGCS: "copy_files_to_gcs.CopyFilesToGCS"
ReadFromDeltaLake: "read_from_delta_lake.ReadFromDeltaLake"
ReadFromDeltaLake: "read_from_delta_lake.ReadFromDeltaLake"
WriteToLakehouse: "write_to_lakehouse.WriteToLakehouse"
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ authors = ["Google Cloud Platform"]
packages = [
{ include = "copy_files_to_gcs.py" },
{ include = "read_from_delta_lake.py" },
{ include = "write_to_lakehouse.py" },
]

[tool.poetry.dependencies]
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Module containing transforms to write data to Lakehouse tables."""

from typing import Iterable, Mapping, Optional
from apache_beam.transforms import PTransform
from apache_beam.yaml.yaml_io import write_to_iceberg


class WriteToLakehouse(PTransform):
"""A PTransform that writes data to a Lakehouse table.

Currently, it wraps the Apache Iceberg sink.
"""

def __init__(
self,
table: str,
catalog_name: Optional[str] = None,
catalog_properties: Optional[Mapping[str, str]] = None,
config_properties: Optional[Mapping[str, str]] = None,
partition_fields: Optional[Iterable[str]] = None,
table_properties: Optional[Mapping[str, str]] = None,
triggering_frequency_seconds: Optional[int] = None,
keep: Optional[Iterable[str]] = None,
drop: Optional[Iterable[str]] = None,
only: Optional[str] = None,
distribution_mode: Optional[str] = None,
autosharding: Optional[bool] = None,
):
super().__init__()
self.table = table
self.catalog_name = catalog_name
self.catalog_properties = catalog_properties
self.config_properties = config_properties
self.partition_fields = partition_fields
self.table_properties = table_properties
self.triggering_frequency_seconds = triggering_frequency_seconds
self.keep = keep
self.drop = drop
self.only = only
self.distribution_mode = distribution_mode
self.autosharding = autosharding

def expand(self, pcoll):
"""Expands the WriteToLakehouse transform."""
return pcoll | write_to_iceberg(
table=self.table,
catalog_name=self.catalog_name,
catalog_properties=self.catalog_properties,
config_properties=self.config_properties,
partition_fields=self.partition_fields,
table_properties=self.table_properties,
triggering_frequency_seconds=self.triggering_frequency_seconds,
keep=self.keep,
drop=self.drop,
only=self.only,
distribution_mode=self.distribution_mode,
autosharding=self.autosharding,
)
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from unittest.mock import MagicMock, patch
from write_to_lakehouse import WriteToLakehouse


class WriteToLakehouseTest(unittest.TestCase):

@patch("write_to_lakehouse.write_to_iceberg")
def test_write_to_lakehouse(self, mock_write_to_iceberg):
mock_transform = MagicMock()
mock_write_to_iceberg.return_value = mock_transform

table = "lakehouse_catalog.dataset.table"
catalog_name = "lakehouse_catalog"
catalog_properties = {"type": "hadoop", "warehouse": "gs://bucket/warehouse"}
config_properties = {"fs.gs.impl": "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem"}
partition_fields = ["day(ts)", "category"]
table_properties = {"commit.retry.num-retries": "2"}
triggering_frequency_seconds = 60
keep = ["field1", "field2"]
drop = ["field3"]
only = "field4"
distribution_mode = "hash"
autosharding = True

transform = WriteToLakehouse(
table=table,
catalog_name=catalog_name,
catalog_properties=catalog_properties,
config_properties=config_properties,
partition_fields=partition_fields,
table_properties=table_properties,
triggering_frequency_seconds=triggering_frequency_seconds,
keep=keep,
drop=drop,
only=only,
distribution_mode=distribution_mode,
autosharding=autosharding,
)

pcoll = MagicMock()
transform.expand(pcoll)

mock_write_to_iceberg.assert_called_once_with(
table=table,
catalog_name=catalog_name,
catalog_properties=catalog_properties,
config_properties=config_properties,
partition_fields=partition_fields,
table_properties=table_properties,
triggering_frequency_seconds=triggering_frequency_seconds,
keep=keep,
drop=drop,
only=only,
distribution_mode=distribution_mode,
autosharding=autosharding,
)


if __name__ == "__main__":
unittest.main()
Loading
Loading