Merge pull request #103 from statisticsnorway/hvr_savings_kolonne

Hvr savings kolonne

Author: SSBHalvor

src/ssb_konjunk/saving.py

@@ -5,6 +5,7 @@
 
 import glob
 import re
+import warnings
 
 import dapla
 import pandas as pd
@@ -350,6 +351,7 @@ def read_ssb_file(
     datatilstand: str = "",
     undermappe: str | None = None,
     filetype: str = "parquet",
+    columns: list[str] | None = None,
     version_number: int | None = None,
     fs: dapla.gcs.GCSFileSystem | None = None,
     seperator: str = ";",
@@ -371,6 +373,7 @@ def read_ssb_file(
         undermappe: Optional folder under 'datatilstand'.
         version_number: possibility to get another version, than the newest (i.e. highest version number). Default: np.nan.
         filetype: the filetype to save as. Default: 'parquet'.
+        columns: Columns to read from the file. If None (default), all columns are read.
         fs: the filesystem, pass with gsc Filesystem if Dapla. Default: None.
         seperator: the seperator to use it filetype is csv. Default: ';'.
         encoding: Encoding for file, base is latin1.
@@ -411,15 +414,27 @@ def read_ssb_file(
         if fs:
             # Samme som tidligere kan brukes til å lese alle filformater.
             with fs.open(file_path, "r") as f:
-                df = pd.read_csv(f, sep=seperator, encoding=encoding)
+                df = pd.read_csv(f, sep=seperator, encoding=encoding, usecols=columns)
                 f.close()
         else:
-            df = pd.read_csv(file_path, sep=seperator, encoding=encoding)
+            df = pd.read_csv(
+                file_path, sep=seperator, encoding=encoding, usecols=columns
+            )
     elif filetype == "parquet":
-        df = pd.read_parquet(file_path, filesystem=fs)
+        df = pd.read_parquet(file_path, columns=columns, filesystem=fs)
     elif filetype == "jsonl":
-        df = pd.read_json(file_path, lines=True)
+        if columns is not None:
+            warnings.warn(
+                f"Columns argumentet blir ignorert for {filetype} filer, hele filen vil bli lastet inn.",
+                stacklevel=2,
+            )
+        df = pd.read_json(file_path, lines=False)
     elif filetype == "json":
+        if columns is not None:
+            warnings.warn(
+                f"Columns argumentet blir ignorert for {filetype} filer, hele filen vil bli lastet inn.",
+                stacklevel=2,
+            )
         df = pd.read_json(file_path, lines=False)
     # Returns pandas df.
     return df