Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/iceberg.py: 22%

1from typing import ( 

2 Any, 

3) 

4 

5from pandas.compat._optional import import_optional_dependency 

6from pandas.util._decorators import set_module 

7 

8from pandas import DataFrame 

9 

10from pandas.io._util import ( 

11 arrow_table_to_pandas, 

12 suppress_pyarrow_values_warning, 

13) 

14 

15 

16@set_module("pandas") 

17def read_iceberg( 

18 table_identifier: str, 

19 catalog_name: str | None = None, 

20 *, 

21 catalog_properties: dict[str, Any] | None = None, 

22 columns: list[str] | None = None, 

23 row_filter: str | None = None, 

24 case_sensitive: bool = True, 

25 snapshot_id: int | None = None, 

26 limit: int | None = None, 

27 scan_properties: dict[str, Any] | None = None, 

28) -> DataFrame: 

29 """ 

30 Read an Apache Iceberg table into a pandas DataFrame. 

31 

32 .. versionadded:: 3.0.0 

33 

34 .. warning:: 

35 

36 read_iceberg is experimental and may change without warning. 

37 

38 Parameters 

39 ---------- 

40 table_identifier : str 

41 Table identifier. 

42 catalog_name : str, optional 

43 The name of the catalog. 

44 catalog_properties : dict of {str: str}, optional 

45 The properties that are used next to the catalog configuration. 

46 columns : list of str, optional 

47 A list of strings representing the column names to return in the output 

48 dataframe. 

49 row_filter : str, optional 

50 A string that describes the desired rows. 

51 case_sensitive : bool, default True 

52 If True column matching is case sensitive. 

53 snapshot_id : int, optional 

54 Snapshot ID to time travel to. By default the table will be scanned as of the 

55 current snapshot ID. 

56 limit : int, optional 

57 An integer representing the number of rows to return in the scan result. 

58 By default all matching rows will be fetched. 

59 scan_properties : dict of {str: obj}, optional 

60 Additional Table properties as a dictionary of string key value pairs to use 

61 for this scan. 

62 

63 Returns 

64 ------- 

65 DataFrame 

66 DataFrame based on the Iceberg table. 

67 

68 See Also 

69 -------- 

70 read_parquet : Read a Parquet file. 

71 

72 Examples 

73 -------- 

74 >>> df = pd.read_iceberg( 

75 ... table_identifier="my_table", 

76 ... catalog_name="my_catalog", 

77 ... catalog_properties={"s3.secret-access-key": "my-secret"}, 

78 ... row_filter="trip_distance >= 10.0", 

79 ... columns=["VendorID", "tpep_pickup_datetime"], 

80 ... ) # doctest: +SKIP 

81 """ 

82 pyiceberg_catalog = import_optional_dependency("pyiceberg.catalog") 

83 pyiceberg_expressions = import_optional_dependency("pyiceberg.expressions") 

84 if catalog_properties is None: 

85 catalog_properties = {} 

86 catalog = pyiceberg_catalog.load_catalog(catalog_name, **catalog_properties) 

87 table = catalog.load_table(table_identifier) 

88 if row_filter is None: 

89 row_filter = pyiceberg_expressions.AlwaysTrue() 

90 if columns is None: 

91 selected_fields = ("*",) 

92 else: 

93 selected_fields = tuple(columns) # type: ignore[assignment] 

94 if scan_properties is None: 

95 scan_properties = {} 

96 result = table.scan( 

97 row_filter=row_filter, 

98 selected_fields=selected_fields, 

99 case_sensitive=case_sensitive, 

100 snapshot_id=snapshot_id, 

101 options=scan_properties, 

102 limit=limit, 

103 ) 

104 return arrow_table_to_pandas(result.to_arrow()) 

105 

106 

107def to_iceberg( 

108 df: DataFrame, 

109 table_identifier: str, 

110 catalog_name: str | None = None, 

111 *, 

112 catalog_properties: dict[str, Any] | None = None, 

113 location: str | None = None, 

114 append: bool = False, 

115 snapshot_properties: dict[str, str] | None = None, 

116) -> None: 

117 """ 

118 Write a DataFrame to an Apache Iceberg table. 

119 

120 .. versionadded:: 3.0.0 

121 

122 Parameters 

123 ---------- 

124 table_identifier : str 

125 Table identifier. 

126 catalog_name : str, optional 

127 The name of the catalog. 

128 catalog_properties : dict of {str: str}, optional 

129 The properties that are used next to the catalog configuration. 

130 location : str, optional 

131 Location for the table. 

132 append : bool, default False 

133 If ``True``, append data to the table, instead of replacing the content. 

134 snapshot_properties : dict of {str: str}, optional 

135 Custom properties to be added to the snapshot summary 

136 

137 See Also 

138 -------- 

139 read_iceberg : Read an Apache Iceberg table. 

140 DataFrame.to_parquet : Write a DataFrame in Parquet format. 

141 """ 

142 pa = import_optional_dependency("pyarrow") 

143 pyiceberg_catalog = import_optional_dependency("pyiceberg.catalog") 

144 if catalog_properties is None: 

145 catalog_properties = {} 

146 catalog = pyiceberg_catalog.load_catalog(catalog_name, **catalog_properties) 

147 with suppress_pyarrow_values_warning(): 

148 arrow_table = pa.Table.from_pandas(df) 

149 table = catalog.create_table_if_not_exists( 

150 identifier=table_identifier, 

151 schema=arrow_table.schema, 

152 location=location, 

153 # we could add `partition_spec`, `sort_order` and `properties` in the 

154 # future, but it may not be trivial without exposing PyIceberg objects 

155 ) 

156 if snapshot_properties is None: 

157 snapshot_properties = {} 

158 if append: 

159 table.append(arrow_table, snapshot_properties=snapshot_properties) 

160 else: 

161 table.overwrite(arrow_table, snapshot_properties=snapshot_properties)