Add union structure family

tiled/_tests/test_writing.py

@@ -19,9 +19,11 @@
 from ..client import Context, from_context, record_history
 from ..queries import Key
 from ..server.app import build_app
-from ..structures.core import Spec
+from ..structures.array import ArrayStructure
+from ..structures.core import Spec, StructureFamily
 from ..structures.data_source import DataSource
 from ..structures.sparse import COOStructure
+from ..structures.table import TableStructure
 from ..validation_registration import ValidationRegistry
 from .utils import fail_with_status_code
 
@@ -451,3 +453,172 @@ async def test_container_export(tree):
         a.write_array([1, 2, 3], key="b")
         buffer = io.BytesIO()
         client.export(buffer, format="application/json")
+
+
+def test_union_one_table(tree):
+    with Context.from_app(build_app(tree)) as context:
+        client = from_context(context)
+        df = pandas.DataFrame({"A": [], "B": []})
+        structure = TableStructure.from_pandas(df)
+        data_source = DataSource(
+            structure_family=StructureFamily.table,
+            structure=structure,
+            name="table",
+        )
+        client.create_union([data_source], key="x")
+
+
+def test_union_two_tables(tree):
+    with Context.from_app(build_app(tree)) as context:
+        client = from_context(context)
+        df1 = pandas.DataFrame({"A": [], "B": []})
+        df2 = pandas.DataFrame({"C": [], "D": [], "E": []})
+        structure1 = TableStructure.from_pandas(df1)
+        structure2 = TableStructure.from_pandas(df2)
+        x = client.create_union(
+            [
+                DataSource(
+                    structure_family=StructureFamily.table,
+                    structure=structure1,
+                    name="table1",
+                ),
+                DataSource(
+                    structure_family=StructureFamily.table,
+                    structure=structure2,
+                    name="table2",
+                ),
+            ],
+            key="x",
+        )
+        x.parts["table1"].write(df1)
+        x.parts["table2"].write(df2)
+        x.parts["table1"].read()
+        x.parts["table2"].read()
+
+
+def test_union_two_tables_colliding_names(tree):
+    with Context.from_app(build_app(tree)) as context:
+        client = from_context(context)
+        df1 = pandas.DataFrame({"A": [], "B": []})
+        df2 = pandas.DataFrame({"C": [], "D": [], "E": []})
+        structure1 = TableStructure.from_pandas(df1)
+        structure2 = TableStructure.from_pandas(df2)
+        with fail_with_status_code(422):
+            client.create_union(
+                [
+                    DataSource(
+                        structure_family=StructureFamily.table,
+                        structure=structure1,
+                        name="table1",
+                    ),
+                    DataSource(
+                        structure_family=StructureFamily.table,
+                        structure=structure2,
+                        name="table1",  # collision
+                    ),
+                ],
+                key="x",
+            )
+
+
+def test_union_two_tables_colliding_keys(tree):
+    with Context.from_app(build_app(tree)) as context:
+        client = from_context(context)
+        df1 = pandas.DataFrame({"A": [], "B": []})
+        df2 = pandas.DataFrame({"A": [], "C": [], "D": []})
+        structure1 = TableStructure.from_pandas(df1)
+        structure2 = TableStructure.from_pandas(df2)
+        with fail_with_status_code(422):
+            client.create_union(
+                [
+                    DataSource(
+                        structure_family=StructureFamily.table,
+                        structure=structure1,
+                        name="table1",
+                    ),
+                    DataSource(
+                        structure_family=StructureFamily.table,
+                        structure=structure2,
+                        name="table2",
+                    ),
+                ],
+                key="x",
+            )
+
+
+def test_union_two_tables_two_arrays(tree):
+    with Context.from_app(build_app(tree)) as context:
+        client = from_context(context)
+        df1 = pandas.DataFrame({"A": [], "B": []})
+        df2 = pandas.DataFrame({"C": [], "D": [], "E": []})
+        arr1 = numpy.ones((5, 5), dtype=numpy.float64)
+        arr2 = 2 * numpy.ones((5, 5), dtype=numpy.int8)
+        structure1 = TableStructure.from_pandas(df1)
+        structure2 = TableStructure.from_pandas(df2)
+        structure3 = ArrayStructure.from_array(arr1)
+        structure4 = ArrayStructure.from_array(arr2)
+        x = client.create_union(
+            [
+                DataSource(
+                    structure_family=StructureFamily.table,
+                    structure=structure1,
+                    name="table1",
+                ),
+                DataSource(
+                    structure_family=StructureFamily.table,
+                    structure=structure2,
+                    name="table2",
+                ),
+                DataSource(
+                    structure_family=StructureFamily.array,
+                    structure=structure3,
+                    name="F",
+                ),
+                DataSource(
+                    structure_family=StructureFamily.array,
+                    structure=structure4,
+                    name="G",
+                ),
+            ],
+            key="x",
+        )
+        # Write by data source.
+        x.parts["table1"].write(df1)
+        x.parts["table2"].write(df2)
+        x.parts["F"].write_block(arr1, (0, 0))
+        x.parts["G"].write_block(arr2, (0, 0))
+
+        # Read by data source.
+        x.parts["table1"].read()
+        x.parts["table2"].read()
+        x.parts["F"].read()
+        x.parts["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:
+        client = from_context(context)
+        df = pandas.DataFrame({"A": [], "B": []})
+        arr = numpy.array([], dtype=numpy.float64)
+        structure1 = TableStructure.from_pandas(df)
+        structure2 = ArrayStructure.from_array(arr)
+        with fail_with_status_code(422):
+            client.create_union(
+                [
+                    DataSource(
+                        structure_family=StructureFamily.table,
+                        structure=structure1,
+                        name="table",
+                    ),
+                    DataSource(
+                        structure_family=StructureFamily.array,
+                        structure=structure2,
+                        name="B",
+                    ),
+                ],
+                key="x",
+            )

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)

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))