Fokko commented on code in PR #5931:
URL: https://github.com/apache/iceberg/pull/5931#discussion_r989937728


##########
python/tests/test_schema.py:
##########
@@ -415,3 +415,133 @@ def test_deserialize_schema(table_schema_simple: Schema):
     )
     expected = table_schema_simple
     assert actual == expected
+
+
+def test_prune_columns_string(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {1}, False) == Schema(
+        NestedField(field_id=1, name="foo", field_type=StringType(), 
required=False), schema_id=1, identifier_field_ids=[1]
+    )
+
+
+def test_prune_columns_string_full(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {1}, True) == Schema(
+        NestedField(field_id=1, name="foo", field_type=StringType(), 
required=False), schema_id=1, identifier_field_ids=[1]
+    )
+
+
+def test_prune_columns_list(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {5}, False) == Schema(
+        NestedField(
+            field_id=4,
+            name="qux",
+            field_type=ListType(type="list", element_id=5, 
element_type=StringType(), element_required=True),
+            required=True,
+        ),
+        schema_id=1,
+        identifier_field_ids=[1],
+    )
+
+
+def test_prune_columns_list_full(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {5}, True) == Schema(
+        NestedField(
+            field_id=4,
+            name="qux",
+            field_type=ListType(type="list", element_id=5, 
element_type=StringType(), element_required=True),
+            required=True,
+        ),
+        schema_id=1,
+        identifier_field_ids=[1],
+    )
+
+
+def test_prune_columns_map(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {9}, False) == Schema(
+        NestedField(
+            field_id=6,
+            name="quux",
+            field_type=MapType(
+                type="map",
+                key_id=7,
+                key_type=StringType(),
+                value_id=8,
+                value_type=MapType(
+                    type="map", key_id=9, key_type=StringType(), value_id=10, 
value_type=IntegerType(), value_required=True
+                ),
+                value_required=True,
+            ),
+            required=True,
+        ),
+        schema_id=1,
+        identifier_field_ids=[1],
+    )
+
+
+def test_prune_columns_map_full(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {9}, True) == Schema(
+        NestedField(
+            field_id=6,
+            name="quux",
+            field_type=MapType(
+                type="map",
+                key_id=7,
+                key_type=StringType(),
+                value_id=8,
+                value_type=MapType(
+                    type="map", key_id=9, key_type=StringType(), value_id=10, 
value_type=IntegerType(), value_required=True
+                ),
+                value_required=True,
+            ),
+            required=True,
+        ),
+        schema_id=1,
+        identifier_field_ids=[1],
+    )
+
+
+def test_prune_columns_map_key(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {10}, False) == Schema(
+        NestedField(
+            field_id=6,
+            name="quux",
+            field_type=MapType(
+                type="map",
+                key_id=7,
+                key_type=StringType(),
+                value_id=8,
+                value_type=MapType(
+                    type="map", key_id=9, key_type=StringType(), value_id=10, 
value_type=IntegerType(), value_required=True
+                ),
+                value_required=True,
+            ),
+            required=True,
+        ),
+        schema_id=1,
+        identifier_field_ids=[1],
+    )
+
+
+def test_prune_columns_struct(table_schema_nested: Schema):
+    assert prune_columns(table_schema_nested, {16}, False) == Schema(
+        NestedField(
+            field_id=15,
+            name="person",
+            field_type=StructType(NestedField(field_id=16, name="name", 
field_type=StringType(), required=False)),
+            required=False,
+        ),
+        schema_id=1,
+        identifier_field_ids=[1],
+    )
+
+
+def test_prune_columns_struct_full(table_schema_nested: Schema):

Review Comment:
   Done 👍🏻 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to