Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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 @@ -3,4 +3,5 @@
packages:
- https://storage.googleapis.com/<package_path>
transforms:
CopyFilesToGCS: "copy_files_to_gcs.CopyFilesToGCS"
CopyFilesToGCS: "copy_files_to_gcs.CopyFilesToGCS"
ReadFromDeltaLake: "read_from_delta_lake.ReadFromDeltaLake"
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ description = "Utility transforms for job builder"
authors = ["Google Cloud Platform"]
packages = [
{ include = "copy_files_to_gcs.py" },
{ include = "read_from_delta_lake.py" },
]

[tool.poetry.dependencies]
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
"""Module containing transforms to read data from Delta Lake tables."""

from typing import Mapping, Optional
from apache_beam.transforms import PTransform
from apache_beam.transforms import managed
from apache_beam.transforms.external import BeamJarExpansionService
from apache_beam.transforms.external import SchemaAwareExternalTransform

DELTA_LAKE_READ_URN = "beam:schematransform:org.apache.beam:delta_lake_read:v1"


class ReadFromDeltaLake(PTransform):
"""A PTransform that reads data from a Delta Lake table.

Args:
table: Identifier or path of the Delta Lake table.
version: Version of the Delta Lake table to read.
timestamp: Timestamp of the Delta Lake table to read.
hadoop_config: Properties passed to Hadoop Configuration.
"""

def __init__(
self,
table: str,
version: Optional[int] = None,
timestamp: Optional[str] = None,
hadoop_config: Optional[Mapping[str, str]] = None,
):
super().__init__()
self.table = table
self.version = version
self.timestamp = timestamp
self.hadoop_config = hadoop_config


def expand(self, pbegin):
"""Expands the ReadFromDeltaLake transform."""
config = {
'table': self.table,
}
if self.version is not None:
config['version'] = self.version
if self.timestamp is not None:
config['timestamp'] = self.timestamp
if self.hadoop_config is not None:
config['hadoop_config'] = dict(self.hadoop_config)

delta_source = getattr(managed, 'DELTA', 'delta')
if hasattr(managed, 'Read') and delta_source in getattr(
managed.Read, '_READ_TRANSFORMS', {}
):
return pbegin | managed.Read(delta_source, config=config)
else:
return pbegin | SchemaAwareExternalTransform(
Comment thread
derrickaw marked this conversation as resolved.
identifier=DELTA_LAKE_READ_URN,
expansion_service=BeamJarExpansionService(
'sdks:java:io:expansion-service:shadowJar'
),
**config,
)


Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
import os
import tempfile
import unittest
from unittest.mock import ANY, MagicMock, patch
import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that, equal_to
from apache_beam.transforms import managed
import pyarrow as pa
import pyarrow.parquet as pq
from read_from_delta_lake import DELTA_LAKE_READ_URN, ReadFromDeltaLake


class ReadFromDeltaLakeTest(unittest.TestCase):
Comment thread
derrickaw marked this conversation as resolved.

@patch('read_from_delta_lake.SchemaAwareExternalTransform')
@patch.object(managed, 'DELTA', 'delta', create=True)
@patch('read_from_delta_lake.managed.Read')
def test_read_from_delta_lake_managed_read(self, mock_managed_read, mock_saet):
mock_managed_read._READ_TRANSFORMS = {'delta': DELTA_LAKE_READ_URN}
mock_transform = MagicMock()
mock_managed_read.return_value = mock_transform


table = 'gs://bucket/delta_table'
version = 2
timestamp = '2026-05-01T12:00:00Z'
hadoop_config = {'fs.gs.project.id': 'test-project'}

transform = ReadFromDeltaLake(
table=table,
version=version,
timestamp=timestamp,
hadoop_config=hadoop_config,
)

pbegin = MagicMock()
transform.expand(pbegin)

mock_managed_read.assert_called_once_with(
'delta',
config={
'table': table,
'version': version,
'timestamp': timestamp,
'hadoop_config': hadoop_config,
},
)


@patch('read_from_delta_lake.SchemaAwareExternalTransform')
@patch('read_from_delta_lake.managed.Read')
def test_read_from_delta_lake_fallback(self, mock_managed_read, mock_saet):
mock_managed_read._READ_TRANSFORMS = {}
mock_transform = MagicMock()
mock_saet.return_value = mock_transform

table = '/path/to/table'
hadoop_config = {'fs.gs.project.id': 'test-project'}

transform = ReadFromDeltaLake(
table=table,
hadoop_config=hadoop_config,
)

pbegin = MagicMock()
transform.expand(pbegin)

mock_saet.assert_called_once_with(
identifier=DELTA_LAKE_READ_URN,
expansion_service=ANY,
table=table,
hadoop_config=hadoop_config,
)

def test_read_from_delta_lake_local_integration(self):
with tempfile.TemporaryDirectory() as temp_dir:
table_dir = os.path.join(temp_dir, 'delta-table')
os.makedirs(table_dir)

# create parquet file
parquet_file_path = os.path.join(table_dir, 'part-00000.parquet')
table = pa.table({'name': ['test-name-1', 'test-name-2']})
pq.write_table(table, parquet_file_path)

file_size = os.path.getsize(parquet_file_path)

log_dir = os.path.join(table_dir, '_delta_log')
os.makedirs(log_dir)

commit_content = (
'{"protocol":{"minReaderVersion":1,"minWriterVersion":2}}\n'
'{"metaData":{"id":"test-id","format":{"provider":"parquet","options":{}},'
'"schemaString":"{\\"type\\":\\"struct\\",\\"fields\\":[{\\"name\\":\\"name\\",\\'
'"type\\":\\"string\\",\\"nullable\\":true,\\"metadata\\":{}}]}",'
'"partitionColumns":[],"configuration":{},"createdAt":123456789}}\n'
f'{{"add":{{"path":"part-00000.parquet","partitionValues":{{}},"size":{file_size},'
'"modificationTime":123456789,"dataChange":true}}\n'
)

# create delta log
with open(os.path.join(log_dir, '00000000000000000000.json'), 'w') as f:
f.write(commit_content)

with TestPipeline() as p:
output = (
p
| ReadFromDeltaLake(table=table_dir)
| beam.Map(lambda row: row._asdict())
)

expected = [{'name': 'test-name-1'}, {'name': 'test-name-2'}]
assert_that(output, equal_to(expected))



if __name__ == '__main__':
unittest.main()


Loading