diff --git a/tiled/_tests/test_writing.py b/tiled/_tests/test_writing.py index 5cbea287f..82b840e29 100644 --- a/tiled/_tests/test_writing.py +++ b/tiled/_tests/test_writing.py @@ -582,15 +582,22 @@ def test_union_two_tables_two_arrays(tree): ], key="x", ) + # Write by data source. x.contents["table1"].write(df1) x.contents["table2"].write(df2) x.contents["F"].write_block(arr1, (0, 0)) x.contents["G"].write_block(arr2, (0, 0)) + + # Read by data source. x.contents["table1"].read() x.contents["table2"].read() x.contents["F"].read() x.contents["G"].read() + # Read by column. + for column in ["A", "B", "C", "D", "E", "F", "G"]: + x[column].read() + def test_union_table_column_array_key_collision(tree): with Context.from_app(build_app(tree)) as context: diff --git a/tiled/adapters/parquet.py b/tiled/adapters/parquet.py index 9c6903bff..1f872e836 100644 --- a/tiled/adapters/parquet.py +++ b/tiled/adapters/parquet.py @@ -74,3 +74,6 @@ def read_partition(self, *args, **kwargs): def structure(self): return self._structure + + def get(self, key): + return self.dataframe_adapter.get(key) diff --git a/tiled/adapters/table.py b/tiled/adapters/table.py index 41b3ce742..6e2937626 100644 --- a/tiled/adapters/table.py +++ b/tiled/adapters/table.py @@ -80,6 +80,11 @@ def __getitem__(self, key): # Must compute to determine shape. return ArrayAdapter.from_array(self.read([key])[key].values) + def get(self, key): + if key not in self.structure().columns: + return None + return ArrayAdapter.from_array(self.read([key])[key].values) + def items(self): yield from ( (key, ArrayAdapter.from_array(self.read([key])[key].values)) diff --git a/tiled/catalog/adapter.py b/tiled/catalog/adapter.py index 0fdc17fe3..4214a3bc5 100644 --- a/tiled/catalog/adapter.py +++ b/tiled/catalog/adapter.py @@ -425,6 +425,10 @@ async def lookup_adapter( for i in range(len(segments)): catalog_adapter = await self.lookup_adapter(segments[:i]) + if (catalog_adapter.structure_family == StructureFamily.union) and len( + segments[i:] + ) == 1: + return await ensure_awaitable(catalog_adapter.get, segments[-1]) if catalog_adapter.data_sources: adapter = await catalog_adapter.get_adapter() for segment in segments[i:]: @@ -969,6 +973,9 @@ class CatalogSparseAdapter(CatalogArrayAdapter): class CatalogTableAdapter(CatalogNodeAdapter): + async def get(self, *args, **kwargs): + return (await self.get_adapter()).get(*args, **kwargs) + async def read(self, *args, **kwargs): return await ensure_awaitable((await self.get_adapter()).read, *args, **kwargs) @@ -987,8 +994,17 @@ async def write_partition(self, *args, **kwargs): class CatalogUnionAdapter(CatalogNodeAdapter): - def get(self, key): - ... + async def get(self, key): + if key not in self.structure().all_keys: + return None + for data_source in self.data_sources: + if data_source.structure_family == StructureFamily.table: + if key in data_source.structure.columns: + return await ensure_awaitable( + self.for_data_source(data_source.name).get, key + ) + if key == data_source.name: + return self.for_data_source(data_source.name) def for_data_source(self, data_source_name): for data_source in self.data_sources: diff --git a/tiled/client/union.py b/tiled/client/union.py index 19cb2ef12..28df62bab 100644 --- a/tiled/client/union.py +++ b/tiled/client/union.py @@ -1,7 +1,7 @@ import copy from .base import STRUCTURE_TYPES, BaseClient -from .utils import client_for_item +from .utils import MSGPACK_MIME_TYPE, ClientError, client_for_item, handle_error class UnionClient(BaseClient): @@ -19,7 +19,31 @@ def contents(self): def __getitem__(self, key): if key not in self.structure().all_keys: raise KeyError(key) - raise NotImplementedError + try: + self_link = self.item["links"]["self"] + if self_link.endswith("/"): + self_link = self_link[:-1] + params = {} + if self._include_data_sources: + params["include_data_sources"] = True + content = handle_error( + self.context.http_client.get( + f"{self_link}/{key}", + headers={"Accept": MSGPACK_MIME_TYPE}, + params=params, + ) + ).json() + except ClientError as err: + if err.response.status_code == 404: + raise KeyError(key) + raise + item = content["data"] + return client_for_item( + self.context, + self.structure_clients, + item, + include_data_sources=self._include_data_sources, + ) class UnionContents: diff --git a/tiled/serialization/table.py b/tiled/serialization/table.py index 8fbead6b8..541375bb2 100644 --- a/tiled/serialization/table.py +++ b/tiled/serialization/table.py @@ -47,7 +47,7 @@ def serialize_csv(df, metadata, preserve_index=False): def deserialize_csv(buffer): import pandas - return pandas.read_csv(io.BytesIO(buffer)) + return pandas.read_csv(io.BytesIO(buffer), headers=False) serialization_registry.register(StructureFamily.table, "text/csv", serialize_csv)