Merge pull request #432 from s22s/fix/431-rf-extract-bits-name-error

Fix rf_local_extract_bits argument parsing

Author: metasim

pyrasterframes/src/main/python/pyrasterframes/rasterfunctions.py

@@ -522,7 +522,7 @@ def rf_local_extract_bits(tile, start_bit, num_bits=1):
     * `startBit` is the first bit to consider, working from the right. It is zero indexed.
     * `numBits` is the number of bits to take moving further to the left. """
     if isinstance(start_bit, int):
-        start_bit = lit(bit_position)
+        start_bit = lit(start_bit)
     if isinstance(num_bits, int):
         num_bits = lit(num_bits)
     return _apply_column_function('rf_local_extract_bits', tile, start_bit, num_bits)

pyrasterframes/src/main/python/tests/RasterFunctionsTests.py

@@ -348,6 +348,16 @@ def test_mask(self):
             expected_data_values
         )
 
+    def test_extract_bits(self):
+        one = np.ones((6, 6), 'uint8')
+        t = Tile(84 * one)
+        df = self.spark.createDataFrame([Row(t=t)])
+        result_py_literals = df.select(rf_local_extract_bits('t', 2, 3)).first()[0]
+        # expect value binary 84 => 1010100 => 101
+        assert_equal(result_py_literals.cells, 5 * one)
+
+        result_cols = df.select(rf_local_extract_bits('t', lit(2), lit(3))).first()[0]
+        assert_equal(result_cols.cells, 5 * one)
 
     def test_resample(self):
         from pyspark.sql.functions import lit