[iceberg] branch master updated: Python: Support visiting MapType (#3989)

[bl...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

blue pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/master by this push:
     new 8748ddd  Python: Support visiting MapType (#3989)
8748ddd is described below

commit 8748ddddd37da1b26d72073858d90cf2965f174c
Author: cccs-eric <er...@cyber.gc.ca>
AuthorDate: Mon Jan 31 11:21:44 2022 -0500

    Python: Support visiting MapType (#3989)
---
 python_legacy/iceberg/api/types/type.py      |  2 +-
 python_legacy/iceberg/api/types/type_util.py | 26 ++++++++++++++++++++++++--
 python_legacy/iceberg/api/types/types.py     |  8 +++++++-
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/python_legacy/iceberg/api/types/type.py b/python_legacy/iceberg/api/types/type.py
index 437b95b..28f000e 100644
--- a/python_legacy/iceberg/api/types/type.py
+++ b/python_legacy/iceberg/api/types/type.py
@@ -63,7 +63,7 @@ class Type(object):
     def as_list_type(self):
         raise ValueError("Not a list type: " + self)
 
-    def asMapType(self):
+    def as_map_type(self):
         raise ValueError("Not a map type: " + self)
 
     def is_nested_type(self):
diff --git a/python_legacy/iceberg/api/types/type_util.py b/python_legacy/iceberg/api/types/type_util.py
index 5a465e5..e0c4428 100644
--- a/python_legacy/iceberg/api/types/type_util.py
+++ b/python_legacy/iceberg/api/types/type_util.py
@@ -143,7 +143,29 @@ def visit(arg, visitor): # noqa: ignore=C901
 
             return visitor.list(list_var, element_result)
         elif type_var.type_id == TypeID.MAP:
-            raise NotImplementedError()
+            map_var = type_var.as_nested_type().as_map_type()
+            visitor.field_ids.append(map_var.key_field.field_id)
+            visitor.field_names.append(map_var.key_field.name)
+            try:
+                key_result = visit(map_var.key_type(), visitor)
+            except NotImplementedError:
+                # will remove it after missing functions are implemented.
+                pass
+            finally:
+                visitor.field_ids.pop()
+                visitor.field_names.pop()
+
+            visitor.field_ids.append(map_var.value_field.field_id)
+            visitor.field_names.append(map_var.value_field.name)
+            try:
+                value_result = visit(map_var.value_type(), visitor)
+            except NotImplementedError:
+                # will remove it after missing functions are implemented.
+                pass
+            finally:
+                visitor.field_ids.pop()
+                visitor.field_names.pop()
+            return visitor.map(map_var, key_result, value_result)
         else:
             return visitor.primitive(arg.as_primitive_type())
     else:
@@ -385,7 +407,7 @@ class IndexById(SchemaVisitor):
             self.index[field.field_id] = field
 
     def map(self, map_var, key_result, value_result):
-        for field in map_var.fields:
+        for field in map_var.fields():
             self.index[field.field_id] = field
 
 
diff --git a/python_legacy/iceberg/api/types/types.py b/python_legacy/iceberg/api/types/types.py
index b5957c2..2bc8d5b 100644
--- a/python_legacy/iceberg/api/types/types.py
+++ b/python_legacy/iceberg/api/types/types.py
@@ -696,6 +696,12 @@ class MapType(NestedType):
     def value_id(self):
         return self.value_field.field_id
 
+    def as_map_type(self):
+        return self
+
+    def is_map_type(self):
+        return True
+
     def is_value_optional(self):
         return self.value_field.is_optional
 
@@ -723,4 +729,4 @@ class MapType(NestedType):
         return MapType.__class__, self.key_field, self.value_field
 
     def _lazy_field_list(self):
-        return tuple(self.key_field, self.value_field)
+        return (self.key_field, self.value_field)